Make the Transformers modeling backend as fast as native vLLM (#47187)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

.buildkite/test-amd.yaml

@@ -472,7 +472,7 @@ steps:
   commands:
   - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
-  - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
+  - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
   - pytest models/language -v -s -m 'distributed(num_gpus=2)'
   - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
   - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
@@ -1637,10 +1637,10 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/models/test_terratorch.py
-  - tests/models/test_transformers.py
+  - tests/models/transformers/test_backend.py
   - tests/models/test_registry.py
   commands:
-  - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+  - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
 
 #-----------------------------------------------------  mi300 · models / language  -----------------------------------------------------#
 
@@ -1860,7 +1860,7 @@ steps:
   - examples/
   commands:
   - pip install --upgrade git+https://github.com/huggingface/transformers
-  - pytest -v -s tests/models/test_transformers.py
+  - pytest -v -s tests/models/transformers/test_backend.py
   - pytest -v -s tests/models/multimodal/test_mapping.py
   - python3 examples/basic/offline_inference/chat.py
   - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl

.buildkite/test_areas/models_basic.yaml

@@ -36,10 +36,10 @@ steps:
   source_file_dependencies:
   - vllm/
   - tests/models/test_terratorch.py
-  - tests/models/test_transformers.py
+  - tests/models/transformers/test_backend.py
   - tests/models/test_registry.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
   mirror:
     amd:
       device: mi325_1
@@ -55,6 +55,7 @@ steps:
   - vllm/
   - tests/models/test_utils.py
   - tests/models/test_vision.py
+  - tests/models/transformers/fusers/
   device: cpu-small
   commands:
-    - pytest -v -s models/test_utils.py models/test_vision.py
+    - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/

.buildkite/test_areas/models_distributed.yaml

@@ -17,7 +17,7 @@ steps:
   - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
   # Avoid importing model tests that cause CUDA reinitialization error
-  - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
+  - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
   - pytest models/language -v -s -m 'distributed(num_gpus=2)'
   - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
   - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py

.github/CODEOWNERS

@@ -119,7 +119,7 @@
 
 # Transformers modeling backend
 /vllm/model_executor/models/transformers @hmellor
-/tests/models/test_transformers.py @hmellor
+/tests/models/transformers @hmellor
 
 # Docs
 /docs/mkdocs @hmellor

docs/models/supported_models.md

@@ -15,7 +15,7 @@ These models are what we list in [supported text models](#list-of-text-only-lang
 
 ### Transformers
 
-vLLM also supports model implementations that are available in Transformers. You should expect the performance of a Transformers model implementation used in vLLM to be within <5% of the performance of a dedicated vLLM model implementation. We call this feature the "Transformers modeling backend".
+vLLM also supports model implementations that are available in Transformers. We call this feature the "Transformers modeling backend". The performance of models loaded with the Transformers modeling backend should be identical to a dedicated vLLM model implementation.
 
 Currently, the Transformers modeling backend works for the following:
 
@@ -140,7 +140,7 @@ Here is what happens in the background when this model is loaded:
 
 That's it!
 
-For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you must add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
+For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you may need to add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
 
 <details class="code">
 <summary>configuration_my_model.py</summary>
@@ -168,9 +168,11 @@ class MyConfig(PretrainedConfig):
 </details>
 
 - `base_model_tp_plan` is a `dict` that maps fully qualified layer name patterns to tensor parallel styles (currently only `"colwise"` and `"rowwise"` are supported).
+    - vLLM infers the tensor parallel style of standard attention (`q`/`k`/`v`/`o_proj`) and gated-MLP/experts (`gate`/`up`/`down_proj`) projections if it can fuse them, so these may not need to be listed. `base_model_tp_plan` is only _required_ for layers that do not follow these patterns; any linear that is neither fused nor named in the plan is replicated.
 - `base_model_pp_plan` is a `dict` that maps direct child layer names to `tuple`s of `list`s of `str`s:
     - You only need to do this for layers which are not present on all pipeline stages
     - vLLM assumes that there will be only one `nn.ModuleList`, which is distributed across the pipeline stages
+    - When no `base_model_pp_plan` is provided, the Transformers modelling backend infers the split from the text model's sole `nn.ModuleList`, keeping the parameter-bearing modules around it (input embeddings, final norm) on the first/last stage (depending on declaration order) and parameter-free modules (e.g. rotary embeddings) on every stage
     - The `list` in the first element of the `tuple` contains the names of the input arguments
     - The `list` in the last element of the `tuple` contains the names of the variables the layer outputs to in your modeling code
 

tests/models/transformers/fusers/test_linear.py

@@ -0,0 +1,480 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Transformers modeling backend's linear fusers."""
+
+import inspect
+from types import MethodType, SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm.model_executor.models.transformers.fuser import get_fuser
+from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser
+
+
+class SiluAndMulStub(nn.Module):
+    """Stand-in for vLLM's `SiluAndMul` (no vLLM config required)."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        return F.silu(x[..., :d]) * x[..., d:]
+
+
+class NoDownGLU(nn.Module):
+    """`act(gate(x)) * up(x)` with no output projection -> `down_name` is None."""
+
+    def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden, inter, bias=bias)
+        self.up_proj = nn.Linear(hidden, inter, bias=bias)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        return self.act_fn(self.gate_proj(x)) * self.up_proj(x)
+
+
+class GLUMLP(NoDownGLU):
+    """`down(act(gate(x)) * up(x))` — the canonical HF GLU MLP."""
+
+    def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False):
+        super().__init__(hidden, inter, bias)
+        self.down_proj = nn.Linear(inter, hidden, bias=bias)
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class ReversedGLUMLP(GLUMLP):
+    """`up(x) * act(gate(x))` — operands swapped (multiply is commutative)."""
+
+    def forward(self, x):
+        return self.down_proj(self.up_proj(x) * self.act_fn(self.gate_proj(x)))
+
+
+class NotAnMLP(nn.Module):
+    """Two linears but no activation*linear multiply -> must not match."""
+
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(8, 8)
+        self.fc2 = nn.Linear(8, 8)
+
+    def forward(self, x):
+        return self.fc2(self.fc1(x))
+
+
+class NotAnActGLUMLP(GLUMLP):
+    """GLU-shaped, but the "activation" is not a known activation module."""
+
+    def __init__(self):
+        super().__init__()
+        self.act_fn = nn.Dropout()
+
+
+class UntraceableMLP(GLUMLP):
+    """Data-dependent control flow *before* the GLU -> no match."""
+
+    def forward(self, x):
+        if x.sum() > 0:  # noqa: SIM108 - intentionally untraceable
+            return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return x
+
+
+class UntraceableTailGLUMLP(GLUMLP):
+    """Data-dependent control flow *after* the GLU -> still fusable."""
+
+    def forward(self, x):
+        y = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        if y.sum() > torch.inf:  # intentionally untraceable
+            y = y * 0
+        return y
+
+
+class FakeAttention(nn.Module):
+    """HF v5-style attention: shape unpacking, dead KV branch, kwargs interface."""
+
+    is_causal = True
+
+    def __init__(
+        self,
+        hidden: int = 32,
+        head_dim: int = 8,
+        heads: int = 4,
+        kv_heads: int = 4,
+        bias: bool = False,
+        layer_idx: int = 0,
+    ):
+        super().__init__()
+        self.config = SimpleNamespace(_attn_implementation="vllm")
+        self.layer_idx = layer_idx
+        self.head_dim = head_dim
+        self.scaling = head_dim**-0.5
+        self.q_proj = nn.Linear(hidden, heads * head_dim, bias=bias)
+        self.k_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias)
+        self.v_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias)
+        self.o_proj = nn.Linear(heads * head_dim, hidden, bias=bias)
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        if past_key_values is not None:
+            k, v = past_key_values.update(k, v, self.layer_idx)
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, attn_weights = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        return self.o_proj(attn_output), attn_weights
+
+
+class ReversedFakeAttention(FakeAttention):
+    """Projections computed in (v, k, q) order — q must still be identified."""
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, None
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
+        )
+        return self.o_proj(attn_output.reshape(*input_shape, -1)), None
+
+
+class ExtraProjAttention(FakeAttention):
+    """A second non-qkv linear of a different width -> `o_proj` still found."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.sink_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+
+
+class QKNormAttention(FakeAttention):
+    """OLMoE-style: a full-dim norm applied to the whole q/k projection output."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.q_norm = nn.RMSNorm(self.q_proj.out_features)
+        self.k_norm = nn.RMSNorm(self.k_proj.out_features)
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        q = self.q_norm(self.q_proj(hidden_states))
+        k = self.k_norm(self.k_proj(hidden_states))
+        v = self.v_proj(hidden_states)
+        return self.o_proj(q + k + v), None
+
+
+class PerHeadQKNormAttention(FakeAttention):
+    """Qwen3-style: a per-head norm (`head_dim`) applied after the head reshape."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.q_norm = nn.RMSNorm(self.head_dim)
+        self.k_norm = nn.RMSNorm(self.head_dim)
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        shape = (*hidden_states.shape[:-1], -1, self.head_dim)
+        q = self.q_norm(self.q_proj(hidden_states).view(shape))
+        k = self.k_norm(self.k_proj(hidden_states).view(shape))
+        v = self.v_proj(hidden_states).view(shape)
+        return self.o_proj((q + k + v).flatten(-2)), None
+
+
+class FakeSelfAttn(nn.Module):
+    """Stand-in for the vLLM `Attention` looked up in `attention_instances`."""
+
+    def __init__(self):
+        super().__init__()
+        self.impl = SimpleNamespace(scale=None)
+
+    def forward(self, q, k, v):
+        # MHA-shaped stub: any deterministic combination of q/k/v will do
+        return q + 2 * k + 3 * v
+
+
+@pytest.fixture(autouse=True)
+def _clear_fuser_cache():
+    get_fuser.cache_clear()
+    yield
+    get_fuser.cache_clear()
+
+
+def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser):
+    """Apply a fuser using plain stand-ins (merged `nn.Linear` + silu AndMul)."""
+    gate = module.get_submodule(fuser.gate_name)
+    up = module.get_submodule(fuser.up_name)
+    merged = nn.Linear(
+        gate.in_features,
+        gate.out_features + up.out_features,
+        bias=gate.bias is not None,
+    )
+    with torch.no_grad():
+        merged.weight.copy_(torch.cat([gate.weight, up.weight], dim=0))
+        if gate.bias is not None:
+            merged.bias.copy_(torch.cat([gate.bias, up.bias], dim=0))
+    setattr(module, fuser.merged_name, merged)
+    setattr(module, fuser.act_name, SiluAndMulStub())
+    delattr(module, fuser.gate_name)
+    delattr(module, fuser.up_name)
+    module.forward = MethodType(fuser.fused_forward, module)
+    return module
+
+
+def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser):
+    """Apply a fuser using a plain merged `nn.Linear` (no TP sharding)."""
+    q, k, v = (
+        module.get_submodule(name)
+        for name in (fuser.q_name, fuser.k_name, fuser.v_name)
+    )
+    merged = nn.Linear(
+        q.in_features,
+        q.out_features + k.out_features + v.out_features,
+        bias=q.bias is not None,
+    )
+    with torch.no_grad():
+        merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0))
+        if q.bias is not None:
+            merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0))
+    merged.split_sizes = [q.out_features, k.out_features, v.out_features]
+    setattr(module, fuser.merged_name, merged)
+    for name in (fuser.q_name, fuser.k_name, fuser.v_name):
+        delattr(module, name)
+    module.forward = MethodType(fuser.fused_forward, module)
+    return module
+
+
+@pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP])
+@pytest.mark.parametrize("bias", [False, True])
+def test_detects_and_rewrites_glu(mlp_cls, bias):
+    with torch.device("meta"):
+        meta = mlp_cls(bias=bias)
+    fuser = get_fuser(meta)
+    assert isinstance(fuser, GLUFuser)
+    assert (
+        fuser.gate_name,
+        fuser.up_name,
+        fuser.act_name,
+        fuser.down_name,
+    ) == ("gate_proj", "up_proj", "act_fn", "down_proj")
+
+    # The rewritten forward references the merged projection instead of the
+    # sources; the rest of the forward is untouched.
+    names = fuser.fused_forward.__code__.co_names
+    assert "gate_up_proj" in names and "act_fn" in names and "down_proj" in names
+    assert not {"gate_proj", "up_proj"} & set(names)
+
+    # Numerics: the fused forward must match the original on a real instance.
+    real = mlp_cls(bias=bias)
+    for p in real.parameters():
+        nn.init.normal_(p, std=0.05)
+    x = torch.randn(4, 16)
+    expected = real(x)
+    fused = _apply_glu_fuser_with_stubs(real, fuser)
+
+    # Fusion is in place: the module keeps its class and other attributes
+    assert fused is real and type(fused) is mlp_cls
+    torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)
+
+
+def test_glu_identifies_down_projection():
+    """The row projection consuming `act(gate(x)) * up(x)` is identified.
+
+    It is forced to `RowParallelLinear` in `update_attrs` so its sharded input
+    matches the column-parallel merged gate/up; `None` when there is no such
+    projection to force (fusion of gate/up still applies)."""
+    with torch.device("meta"):
+        assert get_fuser(GLUMLP()).down_name == "down_proj"
+        assert get_fuser(ReversedGLUMLP()).down_name == "down_proj"
+        assert get_fuser(NoDownGLU()).down_name is None
+
+
+@pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention])
+@pytest.mark.parametrize("kv_heads", [4, 2])
+def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
+    if attn_cls is ReversedFakeAttention and kv_heads == 4:
+        pytest.skip("MHA q/k/v assignment is order-based by design")
+    with torch.device("meta"):
+        meta = attn_cls(kv_heads=kv_heads)
+    fuser = get_fuser(meta)
+    assert isinstance(fuser, QKVFuser)
+    # q (sharded differently under TP) must be identified exactly; k/v may be
+    # swapped for non-canonical compute order, which is numerically consistent
+    # because the weight mapping and the split indices follow the same
+    # assignment.
+    assert fuser.q_name == "q_proj"
+    assert {fuser.k_name, fuser.v_name} == {"k_proj", "v_proj"}
+    assert fuser.o_name == "o_proj"
+
+    # The projections are merged; everything else stays live Python with its
+    # original semantics (branches, kwargs, attribute reads)
+    code = fuser.fused_forward.__code__
+    names = code.co_names
+    assert "qkv_proj" in names and "split_sizes" in names and "o_proj" in names
+    assert not {"q_proj", "k_proj", "v_proj"} & set(names)
+    if attn_cls is FakeAttention:
+        assert "update" in names  # the cache branch survives
+    assert code.co_flags & inspect.CO_VARKEYWORDS  # **kwargs survives
+
+    # Numerics: the fused forward must match the original on a real instance,
+    # with a different layer_idx than the traced instance (kv_heads == heads so
+    # the q/k/v stub combination is shape-compatible).
+    real = attn_cls(kv_heads=4, layer_idx=3)
+    for p in real.parameters():
+        nn.init.normal_(p, std=0.05)
+    x = torch.randn(1, 5, 32)
+    attention_instances = {3: FakeSelfAttn()}
+    expected, _ = real(x, attention_instances=attention_instances)
+    fused = _apply_qkv_fuser_with_stubs(real, fuser)
+
+    # Fusion is in place: the module keeps its class and other attributes
+    assert fused is real and type(fused) is attn_cls
+    assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None
+    out, _ = fused(x, attention_instances=attention_instances)
+    torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
+
+
+def test_qkv_identifies_output_projection():
+    with torch.device("meta"):
+        assert get_fuser(FakeAttention()).o_name == "o_proj"
+        assert get_fuser(ReversedFakeAttention()).o_name == "o_proj"
+        assert get_fuser(ExtraProjAttention()).o_name == "o_proj"
+        # Norm children (q_norm/k_norm) must not disturb o_proj identification.
+        assert get_fuser(QKNormAttention()).o_name == "o_proj"
+        assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj"
+
+
+def test_fuser_is_cached_per_class():
+    with torch.device("meta"):
+        fuser_a = get_fuser(GLUMLP())
+        fuser_b = get_fuser(GLUMLP())
+    assert fuser_a is fuser_b
+    assert GLUMLP in get_fuser.cache
+
+
+@pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP])
+def test_non_matching_modules_return_none(cls):
+    with torch.device("meta"):
+        module = cls()
+    assert get_fuser(module) is None
+
+
+def test_untraceable_tail_still_fuses():
+    with torch.device("meta"):
+        meta = UntraceableTailGLUMLP()
+    fuser = get_fuser(meta)
+    assert isinstance(fuser, GLUFuser)
+
+    # Numerics: the live tail must survive the rewrite
+    real = UntraceableTailGLUMLP()
+    for p in real.parameters():
+        nn.init.normal_(p, std=0.05)
+    x = torch.randn(4, 16)
+    expected = real(x)
+    fused = _apply_glu_fuser_with_stubs(real, fuser)
+    torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)
+
+
+def test_weight_mappings_are_scoped_to_fused_prefixes():
+    from vllm.model_executor.models.utils import WeightsMapper
+
+    with torch.device("meta"):
+        glu_fuser = get_fuser(GLUMLP())
+        qkv_fuser = get_fuser(FakeAttention())
+
+    mapper = WeightsMapper()
+    for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"):
+        mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix))
+    mapper.orig_to_new_stacked.update(
+        qkv_fuser.orig_to_new_stacked("model.layers.0.self_attn")
+    )
+
+    names = [
+        "model.layers.0.mlp.gate_proj.weight",
+        "model.layers.0.mlp.up_proj.weight",
+        "model.layers.1.mlp.gate_proj.weight",
+        "model.layers.0.self_attn.q_proj.weight",
+        "model.layers.0.self_attn.k_proj.weight",
+        "model.layers.0.self_attn.v_proj.weight",
+        # Unfused modules at other prefixes must be left untouched.
+        "model.layers.2.mlp.experts.0.gate_proj.weight",
+        "model.layers.1.self_attn.q_proj.weight",
+    ]
+    # `apply` rewrites the name and stamps the shard id onto each tensor.
+    weights = [(name, torch.empty(0)) for name in names]
+    mapped = list(mapper.apply(weights))
+    mapped_names = [name for name, _ in mapped]
+    shard_ids = [getattr(data, "shard_id", None) for _, data in mapped]
+
+    assert mapped_names == [
+        "model.layers.0.mlp.gate_up_proj.weight",
+        "model.layers.0.mlp.gate_up_proj.weight",
+        "model.layers.1.mlp.gate_up_proj.weight",
+        "model.layers.0.self_attn.qkv_proj.weight",
+        "model.layers.0.self_attn.qkv_proj.weight",
+        "model.layers.0.self_attn.qkv_proj.weight",
+        # Only the exact fused layers are remapped; everything else is untouched.
+        "model.layers.2.mlp.experts.0.gate_proj.weight",
+        "model.layers.1.self_attn.q_proj.weight",
+    ]
+    assert shard_ids == [0, 1, 0, "q", "k", "v", None, None]
+
+    # The fused layers are exposed to the quantization machinery via their
+    # original constituent projection names (what the checkpoint stores).
+    assert glu_fuser.packed_modules_mapping == {
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+    assert qkv_fuser.packed_modules_mapping == {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+    }
+
+
+@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP])
+def test_unfusable_modules_are_not_fused(cls, default_vllm_config):
+    with torch.device("meta"):
+        module = cls()
+    fuser = get_fuser(module)
+    # Either no pattern matches the class, or this instance fails validation
+    # (`recursive_replace` gates fusion and its weight mappings on `validate`)
+    model_config = default_vllm_config.model_config
+    assert fuser is None or not fuser.validate(module, model_config)
+
+
+def test_act_and_mul_derived_from_module(default_vllm_config):
+    from transformers.activations import GELUTanh, SiLUActivation
+
+    from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul
+
+    assert isinstance(GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul)
+    assert isinstance(GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul)
+    gelu_tanh = GLUFuser._get_act_and_mul(GELUTanh())
+    assert isinstance(gelu_tanh, GeluAndMul) and gelu_tanh.approximate == "tanh"
+    gelu = GLUFuser._get_act_and_mul(nn.GELU())
+    assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none"
+    # Not activations at all -> no fusion
+    assert GLUFuser._get_act_and_mul_name(nn.Dropout()) is None
+    assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None
+    with pytest.raises(ValueError, match="No AndMul equivalent"):
+        GLUFuser._get_act_and_mul(nn.Dropout())

tests/models/transformers/fusers/test_moe.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Transformers modeling backend's MoE fuser."""
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm.model_executor.models.transformers.fusers import MoEBlockFuser
+
+from .test_linear import GLUMLP
+
+
+class TopKRouter(nn.Module):
+    """HF v5 top-k router: `linear -> softmax -> topk (-> renorm)`."""
+
+    def __init__(self, num_experts=8, hidden=16, top_k=2, sigmoid=False):
+        super().__init__()
+        self.top_k = top_k
+        self.sigmoid = sigmoid
+        self.weight = nn.Parameter(torch.zeros(num_experts, hidden))
+
+    def forward(self, hidden_states):
+        logits = F.linear(hidden_states, self.weight)
+        scores = torch.sigmoid(logits) if self.sigmoid else F.softmax(logits, dim=-1)
+        value, index = torch.topk(scores, self.top_k, dim=-1)
+        value = value / value.sum(dim=-1, keepdim=True)
+        return logits, value, index
+
+
+class CorrectionRouter(nn.Module):
+    """Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined."""
+
+    def __init__(self, num_experts=8, hidden=16):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(num_experts, hidden))
+        self.register_buffer("e_score_correction_bias", torch.zeros(num_experts))
+
+    def forward(self, hidden_states):
+        logits = F.linear(hidden_states, self.weight)
+        scores = torch.sigmoid(logits) + self.e_score_correction_bias
+        _, index = torch.topk(scores, 2, dim=-1)
+        return logits, scores, index
+
+
+class BiasedRouter(TopKRouter):
+    """A valid top-k router but not `weight`-only (extra `bias` param) -> declined."""
+
+    def __init__(self):
+        super().__init__()
+        self.bias = nn.Parameter(torch.zeros(8))
+
+    def forward(self, hidden_states):
+        logits = F.linear(hidden_states, self.weight) + self.bias
+        scores = F.softmax(logits, dim=-1)
+        value, index = torch.topk(scores, self.top_k, dim=-1)
+        return logits, value, index
+
+
+class DisconnectedRouter(TopKRouter):
+    """linear+softmax+top-k present but top-k ignores the logits -> not a router."""
+
+    def forward(self, hidden_states):
+        logits = F.linear(hidden_states, self.weight)
+        _ = F.softmax(logits, dim=-1)  # scored, but not consumed by top-k
+        value, index = torch.topk(hidden_states, self.top_k, dim=-1)
+        return logits, value, index
+
+
+class MoEExperts(nn.Module):
+    """Packed experts (3D weights); only its name (`experts`) matters here."""
+
+    def __init__(self, num_experts=8, hidden=16, inter=32):
+        super().__init__()
+        self.gate_up_proj = nn.Parameter(torch.zeros(num_experts, 2 * inter, hidden))
+        self.down_proj = nn.Parameter(torch.zeros(num_experts, hidden, inter))
+
+    def forward(self, hidden_states, index, weights):
+        return hidden_states
+
+
+class MoEBlock(nn.Module):
+    """Single-tensor MoE block (Qwen3-style); subclasses override `_shared`."""
+
+    def __init__(self, router_cls=TopKRouter):
+        super().__init__()
+        self.experts = MoEExperts()
+        self.gate = router_cls()
+
+    def _shared(self, x, logits):
+        """The term added to the experts' output (none for a plain block)."""
+        return 0
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        logits, weights, index = self.gate(x)
+        out = self.experts(x, index, weights) + self._shared(x, logits)
+        return out.reshape(hidden_states.shape)
+
+
+class MoEBlockNoShared(MoEBlock):
+    """No shared-expert child but a gate-derived add -> trace skipped, still fuses."""
+
+    def _shared(self, x, logits):
+        return logits.sum()
+
+
+class MoEBlockShared(MoEBlock):
+    """A block with a shared expert and its sigmoid gate (Qwen2-style)."""
+
+    def __init__(self):
+        super().__init__()
+        self.shared_expert = GLUMLP()
+        self.shared_expert_gate = nn.Linear(16, 1, bias=False)
+
+    def _shared(self, x, logits):
+        return torch.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
+
+
+class MoEBlockSharedNoGate(MoEBlock):
+    """A block with an ungated shared expert -> native, shared passed through."""
+
+    def __init__(self):
+        super().__init__()
+        self.shared_expert = GLUMLP()
+
+    def _shared(self, x, logits):
+        return self.shared_expert(x)
+
+
+class MoEBlockTuple(MoEBlock):
+    """A tuple-returning block (gpt-oss-style) -> must decline."""
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        return self.experts(x, index, weights), index
+
+
+class MoEBlockTupleVar(MoEBlock):
+    """Returns a name bound to a tuple, not a literal tuple -> must still decline."""
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        result = self.experts(x, index, weights), index
+        return result
+
+
+class MoEBlockNestedTupleReturn(MoEBlock):
+    """Tuple `return` in a nested helper; block returns one tensor -> still fuses."""
+
+    def forward(self, hidden_states):
+        def keep(a, b):
+            return a, b
+
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        out, _ = keep(self.experts(x, index, weights), index)
+        return out.reshape(hidden_states.shape)
+
+
+class PlainMLP(nn.Module):
+    """A non-GLU FFN: `down(act(up(x)))`, no gating multiply."""
+
+    def __init__(self, hidden: int = 16, inter: int = 32):
+        super().__init__()
+        self.up_proj = nn.Linear(hidden, inter, bias=False)
+        self.down_proj = nn.Linear(inter, hidden, bias=False)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.up_proj(x)))
+
+
+class MoEBlockSharedNonGLU(MoEBlock):
+    """A non-GLU shared expert -> detected by dataflow (no gate/up merge)."""
+
+    def __init__(self):
+        super().__init__()
+        self.shared_expert = PlainMLP()
+
+    def _shared(self, x, logits):
+        return self.shared_expert(x)
+
+
+class MoEBlockUnaccounted(MoEBlock):
+    """A weight-bearing child outside the fused dataflow (pre-router) -> declined."""
+
+    def __init__(self):
+        super().__init__()
+        self.extra = nn.Linear(16, 16, bias=False)
+
+    def forward(self, hidden_states):
+        x = self.extra(hidden_states.reshape(-1, hidden_states.shape[-1]))
+        _, weights, index = self.gate(x)
+        return self.experts(x, index, weights).reshape(hidden_states.shape)
+
+
+class BufferScale(nn.Module):
+    """A stateful child carrying only a buffer (no parameters)."""
+
+    def __init__(self, hidden: int = 16):
+        super().__init__()
+        self.register_buffer("scale", torch.ones(hidden))
+
+    def forward(self, x):
+        return x * self.scale
+
+
+class MoEBlockUnaccountedBuffer(MoEBlockUnaccounted):
+    """Like `MoEBlockUnaccounted`, but the extra child holds only a buffer."""
+
+    def __init__(self):
+        super().__init__()
+        self.extra = BufferScale()
+
+
+@pytest.mark.parametrize("sigmoid", [False, True])
+def test_moe_fuser_detects_router(sigmoid):
+    with torch.device("meta"):
+        block = MoEBlock(lambda: TopKRouter(sigmoid=sigmoid))
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert isinstance(fuser, MoEBlockFuser)
+    assert fuser.gate_name == "gate"
+    assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax")
+    assert fuser.shared_name is None and fuser.shared_gate_name is None
+
+
+def test_moe_fuser_detects_shared_experts():
+    with torch.device("meta"):
+        block = MoEBlockShared()
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert isinstance(fuser, MoEBlockFuser)
+    assert fuser.shared_name == "shared_expert"
+    assert fuser.shared_gate_name == "shared_expert_gate"
+
+
+def test_moe_fuser_skips_shared_detection_without_extra_children():
+    """With only experts and gate, shared-expert detection (and its block trace)
+    is skipped, so a gate-derived add is not misread as a shared expert."""
+    with torch.device("meta"):
+        block = MoEBlockNoShared()
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert isinstance(fuser, MoEBlockFuser)
+    assert fuser.shared_name is None and fuser.shared_gate_name is None
+
+
+def test_moe_fuser_shared_without_gate():
+    with torch.device("meta"):
+        block = MoEBlockSharedNoGate()
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert isinstance(fuser, MoEBlockFuser)
+    assert fuser.shared_name == "shared_expert"
+    assert fuser.shared_gate_name is None
+
+
+def test_moe_fuser_detects_non_glu_shared_expert():
+    with torch.device("meta"):
+        block = MoEBlockSharedNonGLU()
+    fuser = MoEBlockFuser.match(block, "experts")
+    assert isinstance(fuser, MoEBlockFuser)
+    # Recognised by dataflow (added to the experts' output), though not a GLU.
+    assert fuser.shared_name == "shared_expert"
+    assert fuser.shared_gate_name is None
+
+
+@pytest.mark.parametrize(
+    "block_cls",
+    [
+        lambda: MoEBlock(CorrectionRouter),  # score-correction buffer (grouped)
+        lambda: MoEBlock(BiasedRouter),  # router not weight-only (extra param)
+        MoEBlockTuple,  # tuple-returning block (e.g. gpt-oss)
+        MoEBlockTupleVar,  # tuple returned via a name binding, not a literal
+        MoEBlockUnaccounted,  # weight-bearing child outside the fused dataflow
+        MoEBlockUnaccountedBuffer,  # buffer-only child outside the fused dataflow
+    ],
+)
+def test_moe_fuser_declines_unsupported(block_cls):
+    with torch.device("meta"):
+        block = block_cls()
+    assert MoEBlockFuser.match(block, "experts") is None
+
+
+def test_moe_fuser_ignores_nested_returns():
+    """A tuple `return` inside a nested helper must not decline a block whose own
+    forward returns a single tensor."""
+    with torch.device("meta"):
+        block = MoEBlockNestedTupleReturn()
+    assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser)
+
+
+def test_moe_fuser_router_requires_connected_dataflow():
+    """A gate with linear + softmax + top-k present but not wired as a router
+    (top-k selects over the input, not the scored logits) is not detected."""
+    with torch.device("meta"):
+        block = MoEBlock(DisconnectedRouter)
+    assert MoEBlockFuser.match(block, "experts") is None

tests/models/transformers/fusers/test_rms_norm.py

@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Transformers modeling backend's RMSNorm fuser."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm.model_executor.models.transformers.fuser import get_fuser
+from vllm.model_executor.models.transformers.fusers import RMSNormFuser
+
+
+class RMSNorm(nn.Module):
+    """The canonical HF RMSNorm: `weight * x * rsqrt(mean(x**2) + eps)`."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-5, weight: bool = True):
+        super().__init__()
+        if weight:
+            self.weight = nn.Parameter(torch.ones(hidden))
+        self.variance_epsilon = eps
+
+    def _rms(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
+
+    def forward(self, x):
+        return self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
+
+
+class GemmaRMSNorm(RMSNorm):
+    """Zero-centered weight: `(1 + weight) * normalized`."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__(hidden, eps)
+        self.weight = nn.Parameter(torch.zeros(hidden))
+
+    def forward(self, x):
+        return (1.0 + self.weight) * self._rms(x.to(torch.float32)).to(x.dtype)
+
+
+class WeightlessRMSNorm(RMSNorm):
+    """No scale parameter (e.g. Gemma3n `with_scale=False`)."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__(hidden, eps, weight=False)
+
+    def forward(self, x):
+        return self._rms(x.to(torch.float32)).to(x.dtype)
+
+
+class LayerNorm(RMSNorm):
+    """An RMSNorm not named `*RMSNorm`, keeping the input dtype (no upcast)."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__(hidden, eps)
+
+    def forward(self, x):
+        return self.weight * self._rms(x)
+
+
+class NotAnRMSNorm(RMSNorm):
+    """Mean-subtracting LayerNorm-like math -> not an RMSNorm."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__(hidden, eps)
+
+    def forward(self, x):
+        x = x - x.mean(-1, keepdim=True)
+        variance = x.var(-1, keepdim=True)
+        return self.weight * x / torch.sqrt(variance + self.variance_epsilon)
+
+
+class GatedRMSNorm(RMSNorm):
+    """Second input and tail compute -> not an RMSNorm."""
+
+    def forward(self, x, gate=None):
+        normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
+        return normed * F.silu(gate)
+
+
+class GatedFusedRMSNorm(nn.Module):
+    """Same as GatedRMSNorm, but built on the fused `rms_norm` op -> not an RMSNorm."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden))
+        self.eps = eps
+
+    def forward(self, x, gate=None):
+        return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps) * F.silu(gate)
+
+
+class UntraceableGatedRMSNorm(RMSNorm):
+    """Tracer can't see tail compute in forward, but still has a second input (gate)."""
+
+    def forward(self, x, gate=None):
+        normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
+        if gate.sum() > 0:  # untraceable -> partial graph, no visible tail
+            normed = normed * F.silu(gate)
+        return normed
+
+
+@pytest.mark.parametrize(
+    "cls,eps,zero_centered",
+    [
+        (RMSNorm, 1e-5, False),
+        (GemmaRMSNorm, 1e-6, True),
+        (WeightlessRMSNorm, 1e-6, False),
+        (LayerNorm, 1e-6, False),
+        (torch.nn.RMSNorm, 1e-5, False),  # fused `F.rms_norm` op
+    ],
+)
+def test_detects_rms_norm_variants(cls, eps, zero_centered):
+    with torch.device("meta"):
+        fuser = get_fuser(cls(16, eps=eps))
+    assert isinstance(fuser, RMSNormFuser)
+    assert fuser.zero_centered == zero_centered
+
+
+@pytest.mark.parametrize("cls", [NotAnRMSNorm, nn.LayerNorm, nn.SiLU])
+def test_non_rms_norms_are_not_matched(cls):
+    with torch.device("meta"):
+        module = cls(16) if cls is nn.LayerNorm else cls()
+    assert not isinstance(get_fuser(module), RMSNormFuser)
+
+
+@pytest.mark.parametrize(
+    "cls", [GatedRMSNorm, GatedFusedRMSNorm, UntraceableGatedRMSNorm]
+)
+def test_gated_rms_norm_is_not_fused(cls):
+    with torch.device("meta"):
+        assert not isinstance(get_fuser(cls()), RMSNormFuser)
+
+
+@pytest.mark.parametrize(
+    "cls,expected,zero_centered",
+    [
+        (RMSNorm, "RMSNorm", False),
+        (GemmaRMSNorm, "GemmaRMSNorm", True),
+        (WeightlessRMSNorm, "RMSNorm", False),
+    ],
+)
+def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_config):
+    from vllm.model_executor.layers.layernorm import GemmaRMSNorm as VLLMGemmaRMSNorm
+    from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm
+
+    # `default_vllm_config` supplies the config context the CustomOp needs; the
+    # weightless path reads hidden size from the model config, so stub it.
+    model_config = SimpleNamespace(get_hidden_size=lambda: 16)
+    with torch.device("meta"):
+        module = cls()
+        fuser = get_fuser(module)
+        built = fuser.fuse(module, "norm", model_config, None)
+    from vllm.model_executor.models.transformers.fusers.rms_norm import (
+        TPAwareNormMixin,
+    )
+
+    types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm}
+    assert isinstance(built, types_by_name[expected])
+    assert isinstance(built, TPAwareNormMixin)  # fused norms self-correct under TP
+    assert built.variance_epsilon == module.variance_epsilon
+    assert isinstance(built.weight, nn.Parameter) == (
+        getattr(module, "weight", None) is not None
+    )
+
+
+def test_fused_rms_norm_op_default_eps(default_vllm_config):
+    """`torch.nn.RMSNorm` (a single `F.rms_norm` call) matches via the fast path;
+    its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`."""
+    from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm
+
+    with torch.device("meta"):
+        module = torch.nn.RMSNorm(16)  # forward is a single `F.rms_norm` call
+        fuser = get_fuser(module)
+        assert isinstance(fuser, RMSNormFuser)
+        assert not fuser.zero_centered
+        model_config = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32)
+        built = fuser.fuse(module, "norm", model_config, None)
+    assert isinstance(built, VLLMRMSNorm)
+    assert built.variance_epsilon == torch.finfo(torch.float32).eps
+
+
+def test_eps_is_derived_per_instance(default_vllm_config):
+    """Two instances of the same norm class with different eps must fuse to their
+    own eps: the type-cached fuser holds only structure, not this value."""
+    model_config = SimpleNamespace(get_hidden_size=lambda: 16)
+    with torch.device("meta"):
+        for eps in (1e-5, 1e-6):
+            module = RMSNorm(16, eps=eps)
+            built = get_fuser(module).fuse(module, "norm", model_config, None)
+            assert built.variance_epsilon == eps
+
+
+def test_fused_norm_is_gather_capable(default_vllm_config):
+    """Every fused norm is emitted gather-capable, so a norm on a head-sharded
+    projection (OLMoE-style) self-corrects at runtime with no QKV-specific
+    plumbing. A full-width input skips the gather and equals a plain norm."""
+    from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
+    from vllm.model_executor.models.transformers.fusers import rms_norm
+
+    torch.manual_seed(0)
+    x = torch.randn(4, 16)
+    for gathered_cls, plain_cls in [
+        (rms_norm.TPAwareRMSNorm, RMSNorm),
+        (rms_norm.TPAwareGemmaRMSNorm, GemmaRMSNorm),
+    ]:
+        gathered = gathered_cls(hidden_size=16, eps=1e-6)
+        assert isinstance(gathered, rms_norm.TPAwareNormMixin)
+        plain = plain_cls(hidden_size=16, eps=1e-6)
+        with torch.no_grad():
+            weight = torch.randn(16)
+            gathered.weight.copy_(weight)
+            plain.weight.copy_(weight)
+        torch.testing.assert_close(gathered(x), plain(x))
+
+
+def test_gathered_norm_rejects_uneven_sharding(default_vllm_config):
+    """A sharded input (narrower than the full-width weight) that does not tile
+    the weight evenly across ranks is rejected before any collective."""
+    from vllm.model_executor.models.transformers.fusers import rms_norm
+
+    norm = rms_norm.TPAwareRMSNorm(hidden_size=8, eps=1e-6)
+    norm.tp_size = 2  # emulate TP=2 without a real process group
+    with pytest.raises(ValueError, match="does not tile it evenly"):
+        norm(torch.randn(2, 3))  # 3 * 2 != 8

tests/models/transformers/test_backend.py

@@ -6,10 +6,16 @@ from typing import Any
 
 import pytest
 
-from ..conftest import HfRunner, VllmRunner
-from ..utils import multi_gpu_test, prep_prompts
-from .registry import HF_EXAMPLE_MODELS
-from .utils import check_embeddings_close, check_logprobs_close
+from ...conftest import HfRunner, VllmRunner
+from ...utils import multi_gpu_test, prep_prompts
+from ..registry import HF_EXAMPLE_MODELS
+from ..utils import check_embeddings_close, check_logprobs_close
+
+
+@pytest.fixture(scope="function", autouse=True)
+def enable_pickle(monkeypatch):
+    """`LLM.apply_model` requires pickling a function."""
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
 
 
 def get_model(arch: str) -> str:
@@ -18,6 +24,17 @@ def get_model(arch: str) -> str:
     return model_info.default
 
 
+def get_num_fused(model) -> tuple[int, int]:
+    from vllm.model_executor.layers.linear import (
+        MergedColumnParallelLinear,
+        QKVParallelLinear,
+    )
+
+    glu = sum(isinstance(m, MergedColumnParallelLinear) for m in model.modules())
+    qkv = sum(isinstance(m, QKVParallelLinear) for m in model.modules())
+    return glu, qkv
+
+
 def check_implementation(
     runner_ref: type[HfRunner | VllmRunner],
     runner_test: type[VllmRunner],
@@ -25,6 +42,7 @@ def check_implementation(
     model: str,
     kwargs_ref: dict[str, Any] | None = None,
     kwargs_test: dict[str, Any] | None = None,
+    num_fused: tuple[int, int] = (1, 1),
     **kwargs,
 ):
     if kwargs_ref is None:
@@ -41,6 +59,12 @@ def check_implementation(
         model_config = model_test.llm.llm_engine.model_config
         assert model_config.using_transformers_backend()
 
+        num_layers = model_config.hf_config.get_text_config().num_hidden_layers
+        expected_glu, expected_qkv = num_fused
+        for num_glu, num_qkv in model_test.apply_model(get_num_fused):
+            assert num_glu == expected_glu * num_layers
+            assert num_qkv == expected_qkv * num_layers
+
         outputs_test = model_test.generate_greedy_logprobs(*args)
 
     with runner_ref(model, **kwargs_ref) as model_ref:
@@ -58,11 +82,11 @@ def check_implementation(
 
 
 @pytest.mark.parametrize(
-    "model,model_impl",
+    "model,model_impl,num_fused",
     [
-        ("meta-llama/Llama-3.2-1B-Instruct", "transformers"),
-        ("hmellor/Ilama-3.2-1B", "auto"),  # CUSTOM CODE
-        ("allenai/OLMoE-1B-7B-0924", "transformers"),  # MoE
+        ("meta-llama/Llama-3.2-1B-Instruct", "transformers", (1, 1)),
+        ("hmellor/Ilama-3.2-1B", "auto", (1, 1)),  # CUSTOM CODE
+        ("allenai/OLMoE-1B-7B-0924", "transformers", (0, 1)),  # MoE
     ],
 )  # trust_remote_code=True by default
 def test_models(
@@ -71,6 +95,7 @@ def test_models(
     example_prompts: list[str],
     model: str,
     model_impl: str,
+    num_fused: tuple[int, int],
 ) -> None:
     import transformers
     from packaging.version import Version
@@ -84,7 +109,12 @@ def test_models(
         )
 
     check_implementation(
-        hf_runner, vllm_runner, example_prompts, model, model_impl=model_impl
+        hf_runner,
+        vllm_runner,
+        example_prompts,
+        model,
+        num_fused=num_fused,
+        model_impl=model_impl,
     )
 
 

vllm/model_executor/models/transformers/__init__.py

@@ -16,6 +16,10 @@
 # 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
@@ -32,6 +36,36 @@ from vllm.model_executor.models.transformers.pooling import (
 )
 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): ...

vllm/model_executor/models/transformers/base.py

@@ -31,7 +31,6 @@ from transformers.conversion_mapping import (
     WeightRenaming,
     get_model_conversion_mapping,
 )
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config.utils import getattr_iter
@@ -42,6 +41,7 @@ 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,
@@ -51,6 +51,7 @@ from vllm.model_executor.models.interfaces import (
     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,
@@ -58,7 +59,6 @@ from vllm.model_executor.models.transformers.utils import (
     log_replacement,
     replace_conv_class,
     replace_linear_class,
-    replace_rms_norm_class,
 )
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
@@ -74,37 +74,10 @@ if TYPE_CHECKING:
     from transformers import PreTrainedModel
 
     from vllm.config import VllmConfig
-else:
-    PreTrainedModel = object
 
 logger = init_logger(__name__)
 
 
-def vllm_flash_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_flash_attention_forward
-
-
 class Base(
     nn.Module,
     VllmModel,
@@ -141,6 +114,9 @@ class Base(
         """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
@@ -217,7 +193,7 @@ class Base(
         self.text_config._attn_implementation = "vllm"
         self.config.dtype = torch.get_default_dtype()
 
-    def _get_decoder_cls(self, **kwargs: dict) -> type[PreTrainedModel]:
+    def _get_decoder_cls(self, **kwargs: dict) -> type["PreTrainedModel"]:
         """
         Get the decoder class from the model.
 
@@ -236,7 +212,7 @@ class Base(
 
     def _decorate_cls_for_torch_compile(
         self,
-        cls: type[PreTrainedModel],
+        cls: type["PreTrainedModel"],
         dynamic_arg_dims: dict[str, int] | None,
         enable_if: Callable[["VllmConfig"], bool],
         is_encoder: bool,
@@ -356,12 +332,24 @@ class Base(
         if self.pp_group.world_size <= 1:
             return
 
-        if not self.model.supports_pp_plan:
+        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
             )
-            raise ValueError(
-                f"{type(self.model)} does not support pipeline parallel. {tip}"
+            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]:
@@ -376,10 +364,9 @@ class Base(
 
         module_lists = []
         module_list_idx = None
-        pp_plan = list(self.model._pp_plan.keys())
-        for i, name in enumerate(pp_plan):
+        for i, name in enumerate(names):
             # attrgetter in case the module is nested (e.g. "text_model.layers")
-            if isinstance(attrgetter(name)(self.model), nn.ModuleList):
+            if isinstance(attrgetter(name)(module), nn.ModuleList):
                 module_lists.append(name)
                 module_list_idx = i
 
@@ -389,16 +376,16 @@ class Base(
                 "in the base model are not supported yet!"
             )
         if module_list_idx is None:
-            raise ValueError(f"Could not find `ModuleList` in {type(self.model)}")
+            raise ValueError(f"Could not find `ModuleList` in {type(module)}")
 
         # Layers before module list
-        for name in pp_plan[:module_list_idx]:
+        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)(self.model, PPMissingLayer())
+            attrsetter(name)(module, PPMissingLayer())
 
         # Module list
         start_layer, end_layer = get_pp_indices(
@@ -406,41 +393,61 @@ class Base(
             self.pp_group.rank_in_group,
             self.pp_group.world_size,
         )
-        layers_name = pp_plan[module_list_idx]
+        layers_name = names[module_list_idx]
         # attrgetter in case the module is nested (e.g. "text_model.layers")
-        layers = attrgetter(layers_name)(self.model)
+        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 pp_plan[module_list_idx + 1 :]:
+        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)(self.model, PPMissingLayer())
+                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
-        - `*RMSNorm` with vLLM's `RMSNorm`
+        - `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
+        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
             )
-            raise ValueError(
-                f"{type(self.model)} does not support tensor parallel. {tip}"
+            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():
@@ -479,11 +486,16 @@ class Base(
                     )
                 elif isinstance(child_module, (nn.Conv2d, nn.Conv3d)):
                     new_module = replace_conv_class(child_module)
-                elif child_module.__class__.__name__.endswith("RMSNorm"):
-                    new_module = replace_rms_norm_class(
-                        child_module, self.text_config.hidden_size
+                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
                     )
-                else:
+                    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:
@@ -538,7 +550,7 @@ class Base(
                 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_flash_attention_forward
+                # Transformers, it's updated in vllm_attention_forward
                 scale=head_size**-0.5,
                 num_kv_heads=num_kv_heads,
                 cache_config=self.cache_config,

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,220 @@
+# 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(self.qkv.split_sizes, -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}"
+        template = (
+            f"{', '.join(temps)} = {merged}(__arg__).split({merged}.split_sizes, -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,
+        )
+        # The rewritten forward splits the merged projection into this rank's
+        # shard sizes (see `update_forward`)
+        merged.split_sizes = [
+            merged.num_heads * merged.head_size,
+            merged.num_kv_heads * merged.head_size,
+            merged.num_kv_heads * merged.v_head_size,
+        ]
+        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/moe.py

@@ -26,9 +26,11 @@ 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
 
@@ -37,6 +39,8 @@ from .utils import log_replacement
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
 
+logger = init_logger(__name__)
+
 
 @dataclass
 class TransformersMoEState:
@@ -142,14 +146,15 @@ class MoEMixin(MixtureOfExperts):
         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 mlp in self.mlp_layers:
-            mlp.n_local_physical_experts = num_local_physical_experts
-            mlp.n_physical_experts = num_physical_experts
-            mlp.n_redundant_experts = self.num_redundant_experts
-            mlp.experts.update_expert_map()
+        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
@@ -217,10 +222,13 @@ class MoEMixin(MixtureOfExperts):
                 # 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" and (is_modulelist or is_3d):
+                if child_name == experts_name and (is_modulelist or is_3d):
                     # Alias for readability
-                    mlp = module
+                    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():
@@ -230,64 +238,93 @@ class MoEMixin(MixtureOfExperts):
                     # 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 mlp_param_name, _ in mlp.named_parameters():
-                            if "shared_expert" in mlp_param_name:
+                        for moe_block_param_name, _ in moe_block.named_parameters():
+                            if "shared_expert" in moe_block_param_name:
                                 self.num_shared_experts = 1
                                 break
 
-                    # Replace experts module with FusedMoE
-                    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
-                            dist_group = get_ep_group() if is_sp else get_dp_group()
-                            assert sizes[dist_group.rank_in_group] == topk_ids.shape[0]
-                            (topk_ids,) = dist_group.all_gatherv([topk_ids], 0, sizes)
-                        return topk_weights, topk_ids
-
-                    fused_experts = FusedMoE(
+                    kwargs: dict[str, Any] = dict(
                         num_experts=num_experts,
                         top_k=top_k,
                         hidden_size=hidden_size,
                         intermediate_size=intermediate_size,
                         renormalize=renormalize,
-                        # Hard coded because topk happens in Transformers
                         use_grouped_topk=False,
-                        num_expert_group=num_expert_group,
-                        topk_group=topk_group,
                         quant_config=self.quant_config,
                         prefix=qual_name,
                         activation=activation,
                         enable_eplb=enable_eplb,
                         num_redundant_experts=num_redundant_experts,
                         has_bias=has_bias,
-                        custom_routing_function=partial(
-                            custom_routing_function,
-                            moe_state=moe_state,
-                        ),
-                        runner_cls=TransformersMoERunner,
                         routed_experts_cls=TransformersRoutedExperts,
-                        runner_args={"moe_state": moe_state},
                     )
-                    mlp.experts = fused_experts
+                    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(mlp)
+                    self.mlp_layers.append(moe_block)
                     self.moe_layers.append(fused_experts)
                 else:
                     _recursive_replace(child_module, prefix=qual_name)

vllm/model_executor/models/transformers/utils.py

@@ -23,10 +23,8 @@ from typing import TYPE_CHECKING, Literal
 import torch
 from torch import nn
 
-from vllm.config.utils import getattr_iter
 from vllm.logger import init_logger
 from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
-from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     ReplicatedLinear,
@@ -183,45 +181,6 @@ def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv:
     )
 
 
-def replace_rms_norm_class(rms_norm: nn.Module, hidden_size: int) -> RMSNorm:
-    """Replace a Transformers RMSNorm with vLLM's RMSNorm.
-
-    This method assumes:
-    - Weight is stored as `weight`.
-    - Epsilon is stored as `eps` or `variance_epsilon`.
-    - `with_scale` indicates whether the layer has a weight (Gemma3n only).
-    - `var_hidden_size` is only ever used for Intern vision encoder in vLLM
-    and Transformers doesn't appear to have the same concept.
-    """
-    eps = getattr_iter(rms_norm, ("eps", "variance_epsilon"), 1e-6)
-    kwargs = {"hidden_size": hidden_size, "eps": eps}
-    # Update hidden size if weight is available
-    weight_meta = getattr(rms_norm, "weight", None)
-    if weight_meta is not None:
-        kwargs["hidden_size"] = weight_meta.size(0)
-    # Check if weight is all zeros, which indicates GemmaRMSNorm
-    # We must create a new instance because rms_norm is on meta
-    try:
-        with torch.device("cpu"):
-            weight_test = getattr(rms_norm.__class__(1), "weight", None)
-    except Exception:
-        logger.warning(
-            "Failed to determine if RMSNorm weight is centered on zero or one. "
-            "Defaulting to one."
-        )
-        weight_test = None
-    if weight_test is not None and torch.all(weight_test == 0):
-        return GemmaRMSNorm(**kwargs)
-    # Otherwise assume it's a regular RMSNorm
-    kwargs["has_weight"] = getattr(rms_norm, "with_scale", True)
-    if weight_meta is not None:
-        kwargs["dtype"] = weight_meta.dtype
-    else:
-        # No weight, fall back to weightless RMSNorm
-        kwargs["has_weight"] = False
-    return RMSNorm(**kwargs)
-
-
 def recursive_replace_linear(
     model: nn.Module,
     quant_config: "QuantizationConfig | None",