From 925456281eb304fdf2d7c643e81f69b1a01e3416 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Tue, 30 Jun 2026 15:41:20 +0000
Subject: [PATCH 01/35] More sensible name for vllm_attention_forward

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 vllm/model_executor/models/transformers/base.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index bcda62918f34..df3fc72497eb 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -80,7 +80,7 @@
 logger = init_logger(__name__)
 
 
-def vllm_flash_attention_forward(
+def vllm_attention_forward(
     # Transformers args
     module: torch.nn.Module,
     query: torch.Tensor,
@@ -102,7 +102,7 @@ def vllm_flash_attention_forward(
     return self_attn.forward(query, key, value), None
 
 
-ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_flash_attention_forward
+ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward
 
 
 class Base(
@@ -538,7 +538,7 @@ def create_attention_instances(self) -> dict[int, Attention]:
                 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,

From 54dd4070034e3521b951e7fa9aa3f2163b9a3780 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Tue, 30 Jun 2026 17:05:35 +0000
Subject: [PATCH 02/35] Add fusers

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/fusion.py             | 671 ++++++++++++++++++
 1 file changed, 671 insertions(+)
 create mode 100644 vllm/model_executor/models/transformers/fusion.py

diff --git a/vllm/model_executor/models/transformers/fusion.py b/vllm/model_executor/models/transformers/fusion.py
new file mode 100644
index 000000000000..8dca52bd196e
--- /dev/null
+++ b/vllm/model_executor/models/transformers/fusion.py
@@ -0,0 +1,671 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Transformers modeling backend fusion patterns and utilities.
+
+Patterns are detected by tracing each module class once with `torch.fx` and
+matching on the dataflow graph. The rewrite is performed on the forward's
+*source* (AST), replacing only the matched call expressions: the rest of the
+forward stays live Python, so nothing is baked in. The rewritten forward is
+compiled once per class and bound to each fused instance in place.
+"""
+
+import ast
+import inspect
+import operator
+import textwrap
+import types
+from abc import ABC, abstractmethod
+from collections import UserDict
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from cachetools import cached
+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,
+    QKVParallelLinear,
+)
+from vllm.model_executor.models.utils import 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())
+
+
+def _infer_len(node: fx.Node) -> int | None:
+    """Infer a concrete length for a proxy's value by walking its node chain.
+
+    Allows tracing through shape unpacks and `*`-splats like
+    `(*input_shape, -1, head_dim)`, which precede the patterns in HF attention
+    forwards.
+    """
+    # `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 a single `call_module` node, so matching sees the module's
+    own forward structure (HF activation modules are not decomposed into e.g.
+    `sigmoid * x`). `iter` lets tracing proceed through the shape unpacks which
+    precede the patterns (see `_infer_len`); anything else untraceable just
+    ends the trace early, and the partial graph is matched (see `get_fuser`).
+    """
+
+    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 _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 _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 single-argument `self.<name>(...)` call in `funcdef`.
+
+    Raises unless `name` is referenced exactly once, as such a call, so that
+    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 _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
+    )
+
+
+@dataclass
+class BaseFuser(ABC):
+    """A per-module-class fusion plan and the steps to apply it.
+
+    `match` finds the pattern in a traced graph (used only as evidence of the
+    dataflow) and `update_forward` rewrites and compiles the forward's source;
+    both run once per class and the result is cached (see `get_fuser`). `fuse`
+    then applies the fuser to each instance in place (`validate`,
+    `prepare_quant_config`, `update_attrs`, bind `fused_forward`); the instance
+    keeps its class and every attribute not consumed by the fusion.
+    """
+
+    model_config: "ModelConfig"
+    quant_config: "QuantizationConfig"
+
+    merged_name: ClassVar[str]
+    """Attribute name of the merged module created by `update_attrs`."""
+
+    fused_forward: Callable = field(init=False, repr=False)
+    """The compiled rewritten forward, set by `update_forward`."""
+
+    @property
+    @abstractmethod
+    def orig_to_new_stacked(self) -> dict[str, tuple[str, Any]]:
+        """The mapping from original to new stacked modules to
+        be given to the model's WeightsMapper for this fuser."""
+
+    @classmethod
+    @abstractmethod
+    def match(cls, graph: fx.Graph, module: nn.Module, **kwargs) -> "BaseFuser | None":
+        """Match the pattern in `graph`, returning a fuser if found."""
+
+    @abstractmethod
+    def update_forward(self, module: nn.Module) -> None:
+        """Rewrite `type(module)`'s forward source and compile it.
+
+        Raises if the source does not admit the rewrite (fusion is skipped).
+        """
+
+    @abstractmethod
+    def validate(self, module: nn.Module) -> bool:
+        """Whether this fuser can be applied to `module`."""
+
+    @abstractmethod
+    def update_attrs(
+        self,
+        module: nn.Module,
+        prefix: str,
+    ) -> None:
+        """Update `module`'s submodules to match the rewritten forward."""
+
+    def fuse(
+        self,
+        module: nn.Module,
+        prefix: str,
+    ) -> nn.Module:
+        """Apply this fuser to `module` in place, returning it."""
+        if not self.validate(module):
+            return module
+        self.update_attrs(module, prefix)
+        module.forward = types.MethodType(self.fused_forward, module)
+        return module
+
+
+@dataclass
+class GLUFuser(BaseFuser):
+    """Fuser for the GLU pattern `act(gate(x)) * up(x)`."""
+
+    act_name: str
+    gate_name: str
+    up_name: str
+    merged_name: ClassVar[str] = "gate_up_proj"
+
+    @property
+    def orig_to_new_stacked(self) -> dict[str, tuple[str, int]]:
+        return {
+            f".{self.gate_name}": (f".{self.merged_name}", 0),
+            f".{self.up_name}": (f".{self.merged_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] | None:
+        """Search graph for the GLU pattern `act(gate(x)) * up(x)`."""
+        for node in graph.nodes:
+            if (
+                node.op == "call_function"
+                and node.target == operator.mul
+                and len(node.args) == 2
+                and all(isinstance(arg, fx.Node) for arg in node.args)
+            ):
+                a, b = node.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
+        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, **kwargs) -> "GLUFuser | None":
+        if (glu_nodes := cls._get_glu_nodes(graph, module)) is None:
+            return None
+        act_node, gate_node, up_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
+        ):
+            return cls(act_node.target, gate_node.target, up_node.target, **kwargs)
+        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) -> 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,
+    ) -> 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=self.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 projections so their (meta) params are not expected.
+        delattr(module, self.gate_name)
+        delattr(module, self.up_name)
+
+
+@dataclass
+class QKVFuser(BaseFuser):
+    """Fuser for the attention QKV pattern `q(x), k(x), v(x)`."""
+
+    q_name: str
+    k_name: str
+    v_name: str
+    merged_name: ClassVar[str] = "qkv_proj"
+
+    @property
+    def orig_to_new_stacked(self) -> dict[str, tuple[str, str]]:
+        return {
+            f".{self.q_name}": (f".{self.merged_name}", "q"),
+            f".{self.k_name}": (f".{self.merged_name}", "k"),
+            f".{self.v_name}": (f".{self.merged_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, **kwargs) -> "QKVFuser | None":
+        if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None:
+            return None
+        q, k, v = qkv_nodes
+        fuser = cls(q_name=q.target, k_name=k.target, v_name=v.target, **kwargs)
+        if not fuser.validate(module):
+            return None
+        return fuser
+
+    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) -> 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 = self.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 (head_size is None or q.out_features % head_size == 0)
+            and (head_size is None or 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,
+    ) -> None:
+        head_size = self.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=self.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 projections so their (meta) params are not expected.
+        for name in (self.q_name, self.k_name, self.v_name):
+            delattr(module, name)
+
+    def fuse(
+        self,
+        module: nn.Module,
+        prefix: str,
+    ) -> nn.Module:
+        """Apply this fuser to `module` in place, returning it."""
+        if not self.validate(module):
+            return module
+        self.update_attrs(module, prefix)
+        module.forward = types.MethodType(self.fused_forward, module)
+        return module
+
+
+@cached(cache={}, key=type)
+def get_fuser(module: nn.Module, **kwargs) -> BaseFuser | None:
+    """Trace `type(module)` once and return its fuser (cached per module class).
+
+    Returns `None` if the class matches no fusable pattern."""
+    # Currently, all patterns require at least 2 linears, this is a cheap early skip
+    if sum(isinstance(c, nn.Linear) for c in module.children()) < 2:
+        return None
+    tracer = _AllLeafTracer()
+    try:
+        graph = tracer.trace(module)
+    except Exception as exc:
+        # The graph is only evidence for matching, and the patterns appear at
+        # the top of their forwards, so the partial graph traced before the
+        # failure can still be matched.
+        logger.debug("Could not fully trace %s for fusion: %s", type(module), exc)
+        if (graph := getattr(tracer, "graph", None)) is None:
+            return None
+    for fuser_cls in (GLUFuser, QKVFuser):
+        if (fuser := fuser_cls.match(graph, module, **kwargs)) is not None:
+            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
+    return None
+
+
+class Fusers(UserDict):
+    """Mapping from module class to fuser, for all fusable classes in a model."""
+
+    def __init__(self, model: nn.Module, **kwargs):
+        super().__init__({type(m): get_fuser(m, **kwargs) for m in model.modules()})
+
+    def __getitem__(self, m: nn.Module) -> BaseFuser | None:
+        if (fuser := self.data.get(type(m))) is not None and fuser.validate(m):
+            return fuser
+        return None

From 5a672799196c2e5537fae9967b10eff0bf6c61db Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Tue, 30 Jun 2026 17:06:00 +0000
Subject: [PATCH 03/35] Use fusers as part of recursive_replace

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/base.py               | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index df3fc72497eb..8bc9618ea066 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -51,6 +51,7 @@
     SupportsQuant,
 )
 from vllm.model_executor.models.interfaces_base import VllmModel
+from vllm.model_executor.models.transformers.fusion import BaseFuser, Fusers
 from vllm.model_executor.models.transformers.utils import (
     can_enable_torch_compile,
     get_feature_request_tip,
@@ -426,7 +427,10 @@ def recursive_replace(self):
 
         Currently, this replaces:
 
+        - GLUs with a fused `MergedColumnParallelLinear` + `...AndMul`
+        - Attention QKV projections with a fused `QKVParallelLinear` + split
         - `nn.Linear` with vLLM's tensor parallel linear classes
+        - `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d`
         - `*RMSNorm` with vLLM's `RMSNorm`
         """
         tp_plan = self.model.tp_plan
@@ -441,6 +445,18 @@ def recursive_replace(self):
 
         # Prefix the patterns because we always start from `self.model`
         tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()}
+        # Get possible fusers for this model (cached per class, so this is cheap)
+        kwargs = dict(model_config=self.model_config, quant_config=self.quant_config)
+        fusers = Fusers(self.model, **kwargs)
+        orig_to_new_stacked = self.hf_to_vllm_mapper.orig_to_new_stacked
+
+        def update_fuser_mappings(fuser: BaseFuser):
+            # Only update and apply fuser mappings once per fuser
+            if fuser.orig_to_new_stacked.items() <= orig_to_new_stacked.items():
+                return
+            # Fuser weight mappings must be updated before fusion
+            orig_to_new_stacked.update(fuser.orig_to_new_stacked)
+            self._maybe_apply_model_mapping()
 
         def _recursive_replace(module: nn.Module, prefix: str):
             for child_name, child_module in module.named_children():
@@ -483,6 +499,10 @@ def _recursive_replace(module: nn.Module, prefix: str):
                     new_module = replace_rms_norm_class(
                         child_module, self.text_config.hidden_size
                     )
+                elif (fuser := fusers[child_module]) is not None:
+                    update_fuser_mappings(fuser)
+                    new_module = fuser.fuse(child_module, qual_name)
+                    _recursive_replace(new_module, prefix=qual_name)
                 else:
                     _recursive_replace(child_module, prefix=qual_name)
 

From d9ec10e470886a11b465abc05fac10af5593ea61 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Tue, 30 Jun 2026 17:06:21 +0000
Subject: [PATCH 04/35] Check for fusions in general transformers backend tests

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/test_transformers.py | 33 ++++++++++++++++++++++++++-----
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/tests/models/test_transformers.py b/tests/models/test_transformers.py
index eadc3534c378..5f3519552ca7 100644
--- a/tests/models/test_transformers.py
+++ b/tests/models/test_transformers.py
@@ -25,8 +25,19 @@ 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,
 ):
+    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
+
     if kwargs_ref is None:
         kwargs_ref = {}
     if kwargs_test is None:
@@ -41,6 +52,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 +75,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 +88,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 +102,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,
     )
 
 

From b7180c2f1b42d6388b1d9492fb287bc074d5abb7 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Tue, 30 Jun 2026 17:09:06 +0000
Subject: [PATCH 05/35] Fuser tests

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/test_transformers_fusion.py | 414 +++++++++++++++++++++++
 1 file changed, 414 insertions(+)
 create mode 100644 tests/models/test_transformers_fusion.py

diff --git a/tests/models/test_transformers_fusion.py b/tests/models/test_transformers_fusion.py
new file mode 100644
index 000000000000..fec14acc1dc8
--- /dev/null
+++ b/tests/models/test_transformers_fusion.py
@@ -0,0 +1,414 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Transformers backend's fx-based fusion (GLU and QKV).
+
+These exercise the structural detection and surgical subgraph rewrite in
+`vllm.model_executor.models.transformers.fusion` without needing a distributed
+environment; the real `MergedColumnParallelLinear` / `QKVParallelLinear` /
+`...AndMul` integration is covered end-to-end by `test_fusion`, `test_models`
+and `test_quantization` in `tests/models/test_transformers.py`.
+"""
+
+import inspect
+import types
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Registers the "vllm" attention interface in ALL_ATTENTION_FUNCTIONS
+import vllm.model_executor.models.transformers.base  # noqa: F401
+from vllm.model_executor.models.transformers import fusion
+from vllm.model_executor.models.transformers.fusion import (
+    GLUFuser,
+    QKVFuser,
+    get_fuser,
+)
+
+
+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 GLUMLP(nn.Module):
+    """`act(gate(x)) * up(x)` — the canonical HF GLU MLP."""
+
+    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.down_proj = nn.Linear(inter, hidden, bias=bias)
+        self.act_fn = nn.SiLU()
+
+    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 pattern -> no match (the
+    partial trace ends before reaching the GLU)."""
+
+    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 pattern -> still fusable (the
+    partial trace contains the GLU, and the AST rewrite keeps the tail live)."""
+
+    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):
+    """Mimics the structure of a HF v5 attention module under vLLM.
+
+    The forward exercises the patterns the fusion tracer must support: shape
+    unpacking, the dead KV cache branch, and the `attention_interface` call
+    splatting `**kwargs` (resolved to vLLM's interface, which is traced
+    through).
+    """
+
+    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 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():
+    fusion.get_fuser.cache_clear()
+    yield
+    fusion.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 = types.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 = types.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,
+    ) == ("gate_proj", "up_proj", "act_fn")
+
+    # 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)
+
+
+@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"}
+
+    # 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_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 fusion.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_substr.update(glu_fuser.weight_mappings(prefix))
+    mapper.orig_to_new_substr.update(
+        qkv_fuser.weight_mappings("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",
+    ]
+    mapped = mapper.apply_list(names)
+    assert mapped == [
+        "model.layers.0.mlp.gate_up_proj.0.weight",
+        "model.layers.0.mlp.gate_up_proj.1.weight",
+        "model.layers.1.mlp.gate_up_proj.0.weight",
+        "model.layers.0.self_attn.qkv_proj.q.weight",
+        "model.layers.0.self_attn.qkv_proj.k.weight",
+        "model.layers.0.self_attn.qkv_proj.v.weight",
+        "model.layers.2.mlp.experts.0.gate_proj.weight",
+        "model.layers.1.self_attn.q_proj.weight",
+    ]
+    # The mappings must also be visible to the quantization machinery,
+    # which derives packed modules from substr (not regex) entries.
+    assert mapper.get_packed_modules_mapping() == {
+        "gate_up_proj": ["gate_up_proj.0", "gate_up_proj.1"],
+        "qkv_proj": ["qkv_proj.q", "qkv_proj.k", "qkv_proj.v"],
+    }
+
+
+@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP])
+def test_unfusable_modules_are_not_fused(cls):
+    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`)
+    assert fuser is None or not fuser.validate(module)
+
+
+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(fusion.GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul)
+    assert isinstance(fusion.GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul)
+    gelu_tanh = fusion.GLUFuser._get_act_and_mul(GELUTanh())
+    assert isinstance(gelu_tanh, GeluAndMul) and gelu_tanh.approximate == "tanh"
+    gelu = fusion.GLUFuser._get_act_and_mul(nn.GELU())
+    assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none"
+    # Not activations at all -> no fusion
+    assert fusion.GLUFuser._get_act_and_mul_name(nn.Dropout()) is None
+    assert fusion.GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None
+    with pytest.raises(ValueError, match="No AndMul equivalent"):
+        fusion.GLUFuser._get_act_and_mul(nn.Dropout())

From 4293df9419635afe45e385f5298c30f2b0a00913 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Wed, 1 Jul 2026 09:15:25 +0000
Subject: [PATCH 06/35] rework a bit

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/test_transformers_fusion.py      |  47 +++--
 .../models/transformers/base.py               |  33 ++--
 .../models/transformers/fusion.py             | 180 +++++++++---------
 3 files changed, 136 insertions(+), 124 deletions(-)

diff --git a/tests/models/test_transformers_fusion.py b/tests/models/test_transformers_fusion.py
index fec14acc1dc8..76e93489e094 100644
--- a/tests/models/test_transformers_fusion.py
+++ b/tests/models/test_transformers_fusion.py
@@ -351,9 +351,9 @@ def test_weight_mappings_are_scoped_to_fused_prefixes():
 
     mapper = WeightsMapper()
     for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"):
-        mapper.orig_to_new_substr.update(glu_fuser.weight_mappings(prefix))
-    mapper.orig_to_new_substr.update(
-        qkv_fuser.weight_mappings("model.layers.0.self_attn")
+        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 = [
@@ -367,33 +367,44 @@ def test_weight_mappings_are_scoped_to_fused_prefixes():
         "model.layers.2.mlp.experts.0.gate_proj.weight",
         "model.layers.1.self_attn.q_proj.weight",
     ]
-    mapped = mapper.apply_list(names)
-    assert mapped == [
-        "model.layers.0.mlp.gate_up_proj.0.weight",
-        "model.layers.0.mlp.gate_up_proj.1.weight",
-        "model.layers.1.mlp.gate_up_proj.0.weight",
-        "model.layers.0.self_attn.qkv_proj.q.weight",
-        "model.layers.0.self_attn.qkv_proj.k.weight",
-        "model.layers.0.self_attn.qkv_proj.v.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",
     ]
-    # The mappings must also be visible to the quantization machinery,
-    # which derives packed modules from substr (not regex) entries.
-    assert mapper.get_packed_modules_mapping() == {
-        "gate_up_proj": ["gate_up_proj.0", "gate_up_proj.1"],
-        "qkv_proj": ["qkv_proj.q", "qkv_proj.k", "qkv_proj.v"],
+    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):
+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`)
-    assert fuser is None or not fuser.validate(module)
+    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):
diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 8bc9618ea066..9a9068c07421 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -142,6 +142,9 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
         """Ignore unexpected weights whose qualname starts with these prefixes."""
         self.ignore_unexpected_suffixes: list[str] = []
         """Ignore unexpected weights whose qualname ends with these suffixes."""
+        self.packed_modules_mapping: dict[str, list[str]] = {}
+        """Fused module -> constituent projections, populated by `recursive_replace`
+        for the quantization machinery and loaders (e.g. bitsandbytes)."""
 
         # Attrs for Eagle3 (see self.set_aux_hidden_state_layers)
         self._target_class: type[nn.Module] = nn.Module
@@ -445,18 +448,18 @@ def recursive_replace(self):
 
         # Prefix the patterns because we always start from `self.model`
         tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()}
-        # Get possible fusers for this model (cached per class, so this is cheap)
-        kwargs = dict(model_config=self.model_config, quant_config=self.quant_config)
-        fusers = Fusers(self.model, **kwargs)
-        orig_to_new_stacked = self.hf_to_vllm_mapper.orig_to_new_stacked
-
-        def update_fuser_mappings(fuser: BaseFuser):
-            # Only update and apply fuser mappings once per fuser
-            if fuser.orig_to_new_stacked.items() <= orig_to_new_stacked.items():
-                return
-            # Fuser weight mappings must be updated before fusion
-            orig_to_new_stacked.update(fuser.orig_to_new_stacked)
-            self._maybe_apply_model_mapping()
+        # 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():
@@ -500,8 +503,10 @@ def _recursive_replace(module: nn.Module, prefix: str):
                         child_module, self.text_config.hidden_size
                     )
                 elif (fuser := fusers[child_module]) is not None:
-                    update_fuser_mappings(fuser)
-                    new_module = fuser.fuse(child_module, qual_name)
+                    register_fusion(fuser, qual_name)
+                    new_module = fuser.fuse(
+                        child_module, qual_name, self.model_config, self.quant_config
+                    )
                     _recursive_replace(new_module, prefix=qual_name)
                 else:
                     _recursive_replace(child_module, prefix=qual_name)
diff --git a/vllm/model_executor/models/transformers/fusion.py b/vllm/model_executor/models/transformers/fusion.py
index 8dca52bd196e..38c7b2da4c4c 100644
--- a/vllm/model_executor/models/transformers/fusion.py
+++ b/vllm/model_executor/models/transformers/fusion.py
@@ -1,12 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Transformers modeling backend fusion patterns and utilities.
+"""Fusion patterns for the Transformers modeling backend.
 
-Patterns are detected by tracing each module class once with `torch.fx` and
-matching on the dataflow graph. The rewrite is performed on the forward's
-*source* (AST), replacing only the matched call expressions: the rest of the
-forward stays live Python, so nothing is baked in. The rewritten forward is
-compiled once per class and bound to each fused instance in place.
+Each module class is traced once with `torch.fx` to detect a pattern from its
+dataflow, then its forward *source* is rewritten (AST) to replace only the
+matched calls; the rest stays live Python. The rewritten forward is compiled
+once per class and bound to each fused instance in place.
 """
 
 import ast
@@ -18,7 +17,7 @@
 from collections import UserDict
 from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
 from cachetools import cached
 from torch import fx, nn
@@ -33,7 +32,7 @@
     MergedColumnParallelLinear,
     QKVParallelLinear,
 )
-from vllm.model_executor.models.utils import maybe_prefix
+from vllm.model_executor.models.utils import ShardId, maybe_prefix
 
 if TYPE_CHECKING:
     from vllm.config.model import ModelConfig
@@ -51,11 +50,10 @@
 
 
 def _infer_len(node: fx.Node) -> int | None:
-    """Infer a concrete length for a proxy's value by walking its node chain.
+    """Concrete length of a proxy's value, inferred from its node chain.
 
-    Allows tracing through shape unpacks and `*`-splats like
-    `(*input_shape, -1, head_dim)`, which precede the patterns in HF attention
-    forwards.
+    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 (
@@ -95,11 +93,10 @@ def __len__(self) -> int:
 class _AllLeafTracer(fx.Tracer):
     """Tracer that treats every submodule as a leaf.
 
-    Each child stays a single `call_module` node, so matching sees the module's
-    own forward structure (HF activation modules are not decomposed into e.g.
-    `sigmoid * x`). `iter` lets tracing proceed through the shape unpacks which
-    precede the patterns (see `_infer_len`); anything else untraceable just
-    ends the trace early, and the partial graph is matched (see `get_fuser`).
+    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:
@@ -157,10 +154,10 @@ def _compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable:
 
 
 def _single_self_call(funcdef: ast.FunctionDef, name: str) -> ast.Call:
-    """The unique single-argument `self.<name>(...)` call in `funcdef`.
+    """The unique `self.<name>(arg)` call in `funcdef`.
 
-    Raises unless `name` is referenced exactly once, as such a call, so that
-    the source rewrite agrees with the fx match.
+    Raises unless `name` appears exactly once, as such a call, so the source
+    rewrite agrees with the fx match.
     """
     uses = [
         node
@@ -231,18 +228,13 @@ def _is_linear(node: fx.Node, module: nn.Module) -> bool:
 
 @dataclass
 class BaseFuser(ABC):
-    """A per-module-class fusion plan and the steps to apply it.
-
-    `match` finds the pattern in a traced graph (used only as evidence of the
-    dataflow) and `update_forward` rewrites and compiles the forward's source;
-    both run once per class and the result is cached (see `get_fuser`). `fuse`
-    then applies the fuser to each instance in place (`validate`,
-    `prepare_quant_config`, `update_attrs`, bind `fused_forward`); the instance
-    keeps its class and every attribute not consumed by the fusion.
-    """
+    """A per-class fusion plan and the steps to apply it.
 
-    model_config: "ModelConfig"
-    quant_config: "QuantizationConfig"
+    `match` and `update_forward` analyse the module *class* once (cached, see
+    `get_fuser`); `fuse` then 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`."""
@@ -252,43 +244,65 @@ class BaseFuser(ABC):
 
     @property
     @abstractmethod
-    def orig_to_new_stacked(self) -> dict[str, tuple[str, Any]]:
-        """The mapping from original to new stacked modules to
-        be given to the model's WeightsMapper for this fuser."""
+    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]}
 
     @classmethod
     @abstractmethod
-    def match(cls, graph: fx.Graph, module: nn.Module, **kwargs) -> "BaseFuser | None":
+    def match(cls, graph: fx.Graph, module: nn.Module) -> "BaseFuser | None":
         """Match the pattern in `graph`, returning a fuser if found."""
 
     @abstractmethod
     def update_forward(self, module: nn.Module) -> None:
-        """Rewrite `type(module)`'s forward source and compile it.
+        """Rewrite and compile `type(module)`'s forward source.
 
-        Raises if the source does not admit the rewrite (fusion is skipped).
+        Raises if the source does not admit the rewrite (fusion is then skipped).
         """
 
     @abstractmethod
-    def validate(self, module: nn.Module) -> bool:
-        """Whether this fuser can be applied to `module`."""
+    def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
+        """Whether this fuser can be applied to this `module` instance."""
 
     @abstractmethod
     def update_attrs(
         self,
         module: nn.Module,
         prefix: str,
+        model_config: "ModelConfig",
+        quant_config: "QuantizationConfig",
     ) -> None:
-        """Update `module`'s submodules to match the rewritten forward."""
+        """Replace `module`'s submodules with the merged module."""
 
     def fuse(
         self,
         module: nn.Module,
         prefix: str,
+        model_config: "ModelConfig",
+        quant_config: "QuantizationConfig",
     ) -> nn.Module:
-        """Apply this fuser to `module` in place, returning it."""
-        if not self.validate(module):
-            return module
-        self.update_attrs(module, prefix)
+        """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
 
@@ -303,11 +317,8 @@ class GLUFuser(BaseFuser):
     merged_name: ClassVar[str] = "gate_up_proj"
 
     @property
-    def orig_to_new_stacked(self) -> dict[str, tuple[str, int]]:
-        return {
-            f".{self.gate_name}": (f".{self.merged_name}", 0),
-            f".{self.up_name}": (f".{self.merged_name}", 1),
-        }
+    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:
@@ -366,7 +377,7 @@ def _get_act_and_mul(cls, act: nn.Module) -> nn.Module:
         raise ValueError(f"No AndMul equivalent for {type(act)}")
 
     @classmethod
-    def match(cls, graph: fx.Graph, module: nn.Module, **kwargs) -> "GLUFuser | None":
+    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 = glu_nodes
@@ -377,7 +388,7 @@ def match(cls, graph: fx.Graph, module: nn.Module, **kwargs) -> "GLUFuser | None
         if gate.in_features == up.in_features and (gate.bias is None) == (
             up.bias is None
         ):
-            return cls(act_node.target, gate_node.target, up_node.target, **kwargs)
+            return cls(act_node.target, gate_node.target, up_node.target)
         return None
 
     def update_forward(self, module: nn.Module) -> None:
@@ -406,7 +417,7 @@ def update_forward(self, module: nn.Module) -> None:
         _replace_expr(funcdef, muls[0], act_call)
         self.fused_forward = _compile_forward(funcdef, fn)
 
-    def validate(self, module: nn.Module) -> bool:
+    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))
@@ -417,6 +428,8 @@ 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)
@@ -425,7 +438,7 @@ def update_attrs(
             input_size=gate.in_features,
             output_sizes=[gate.out_features, up.out_features],
             bias=gate.bias is not None,
-            quant_config=self.quant_config,
+            quant_config=quant_config,
             prefix=maybe_prefix(prefix, self.merged_name),
             return_bias=False,
         )
@@ -440,7 +453,7 @@ def update_attrs(
         )
         setattr(module, self.merged_name, merged)
         setattr(module, self.act_name, act_fn)
-        # Drop the consumed projections so their (meta) params are not expected.
+        # Drop the consumed submodules so their (meta) params are not expected.
         delattr(module, self.gate_name)
         delattr(module, self.up_name)
 
@@ -455,12 +468,8 @@ class QKVFuser(BaseFuser):
     merged_name: ClassVar[str] = "qkv_proj"
 
     @property
-    def orig_to_new_stacked(self) -> dict[str, tuple[str, str]]:
-        return {
-            f".{self.q_name}": (f".{self.merged_name}", "q"),
-            f".{self.k_name}": (f".{self.merged_name}", "k"),
-            f".{self.v_name}": (f".{self.merged_name}", "v"),
-        }
+    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(
@@ -494,14 +503,11 @@ def _get_qkv_nodes(
         return q_node, k_node, v_node
 
     @classmethod
-    def match(cls, graph: fx.Graph, module: nn.Module, **kwargs) -> "QKVFuser | None":
+    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
-        fuser = cls(q_name=q.target, k_name=k.target, v_name=v.target, **kwargs)
-        if not fuser.validate(module):
-            return None
-        return fuser
+        return cls(q_name=q.target, k_name=k.target, v_name=v.target)
 
     def update_forward(self, module: nn.Module) -> None:
         """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source."""
@@ -558,18 +564,18 @@ def update_forward(self, module: nn.Module) -> None:
             _replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))
         self.fused_forward = _compile_forward(funcdef, fn)
 
-    def validate(self, module: nn.Module) -> bool:
+    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 = self.model_config.get_head_size()
+        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 (head_size is None or q.out_features % head_size == 0)
-            and (head_size is None or k.out_features % head_size == 0)
+            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))
@@ -579,8 +585,10 @@ def update_attrs(
         self,
         module: nn.Module,
         prefix: str,
+        model_config: "ModelConfig",
+        quant_config: "QuantizationConfig",
     ) -> None:
-        head_size = self.model_config.get_head_size()
+        head_size = model_config.get_head_size()
         q = module.get_submodule(self.q_name)
         k = module.get_submodule(self.k_name)
         merged = QKVParallelLinear(
@@ -589,7 +597,7 @@ def update_attrs(
             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=self.quant_config,
+            quant_config=quant_config,
             prefix=maybe_prefix(prefix, self.merged_name),
             return_bias=False,
         )
@@ -612,28 +620,14 @@ def update_attrs(
             merged.num_kv_heads * merged.v_head_size,
         ]
         setattr(module, self.merged_name, merged)
-        # Drop the consumed projections so their (meta) params are not expected.
+        # 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)
 
-    def fuse(
-        self,
-        module: nn.Module,
-        prefix: str,
-    ) -> nn.Module:
-        """Apply this fuser to `module` in place, returning it."""
-        if not self.validate(module):
-            return module
-        self.update_attrs(module, prefix)
-        module.forward = types.MethodType(self.fused_forward, module)
-        return module
-
 
 @cached(cache={}, key=type)
-def get_fuser(module: nn.Module, **kwargs) -> BaseFuser | None:
-    """Trace `type(module)` once and return its fuser (cached per module class).
-
-    Returns `None` if the class matches no fusable pattern."""
+def get_fuser(module: nn.Module) -> BaseFuser | None:
+    """The fuser for `type(module)` (cached per class), or `None` if no match."""
     # Currently, all patterns require at least 2 linears, this is a cheap early skip
     if sum(isinstance(c, nn.Linear) for c in module.children()) < 2:
         return None
@@ -648,7 +642,7 @@ def get_fuser(module: nn.Module, **kwargs) -> BaseFuser | None:
         if (graph := getattr(tracer, "graph", None)) is None:
             return None
     for fuser_cls in (GLUFuser, QKVFuser):
-        if (fuser := fuser_cls.match(graph, module, **kwargs)) is not None:
+        if (fuser := fuser_cls.match(graph, module)) is not None:
             try:
                 fuser.update_forward(module)
             except Exception as exc:
@@ -662,10 +656,12 @@ def get_fuser(module: nn.Module, **kwargs) -> BaseFuser | None:
 class Fusers(UserDict):
     """Mapping from module class to fuser, for all fusable classes in a model."""
 
-    def __init__(self, model: nn.Module, **kwargs):
-        super().__init__({type(m): get_fuser(m, **kwargs) for m in model.modules()})
+    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:
-        if (fuser := self.data.get(type(m))) is not None and fuser.validate(m):
+        fuser = self.data.get(type(m))
+        if fuser is not None and fuser.validate(m, self.model_config):
             return fuser
         return None

From 2772c102e9ae978943d452d94ae24efa137b7239 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Wed, 1 Jul 2026 12:27:19 +0000
Subject: [PATCH 07/35] Actually run the new fusion tests

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .buildkite/test_areas/models_basic.yaml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml
index 5eb799efa187..8491987f964d 100644
--- a/.buildkite/test_areas/models_basic.yaml
+++ b/.buildkite/test_areas/models_basic.yaml
@@ -60,6 +60,7 @@ steps:
   - vllm/
   - tests/models/test_utils.py
   - tests/models/test_vision.py
+  - tests/models/test_transformers_fusion.py
   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/test_transformers_fusion.py

From 6ef790d10628a597a2d4a85ec8125fc8ad540323 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:47:15 +0000
Subject: [PATCH 08/35] Refactor so each fuser gets its own file and fx utils
 are separated

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/base.py               |   2 +-
 .../models/transformers/fuser.py              |  78 ++
 .../models/transformers/fusers/__init__.py    |  10 +
 .../models/transformers/fusers/base.py        | 130 ++++
 .../models/transformers/fusers/glu.py         | 193 +++++
 .../models/transformers/fusers/qkv.py         | 195 +++++
 .../models/transformers/fusers/rms_norm.py    | 139 ++++
 .../models/transformers/fusion.py             | 667 ------------------
 .../models/transformers/fx_utils.py           | 239 +++++++
 9 files changed, 985 insertions(+), 668 deletions(-)
 create mode 100644 vllm/model_executor/models/transformers/fuser.py
 create mode 100644 vllm/model_executor/models/transformers/fusers/__init__.py
 create mode 100644 vllm/model_executor/models/transformers/fusers/base.py
 create mode 100644 vllm/model_executor/models/transformers/fusers/glu.py
 create mode 100644 vllm/model_executor/models/transformers/fusers/qkv.py
 create mode 100644 vllm/model_executor/models/transformers/fusers/rms_norm.py
 delete mode 100644 vllm/model_executor/models/transformers/fusion.py
 create mode 100644 vllm/model_executor/models/transformers/fx_utils.py

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 9a9068c07421..5eac3eed3320 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -51,7 +51,7 @@
     SupportsQuant,
 )
 from vllm.model_executor.models.interfaces_base import VllmModel
-from vllm.model_executor.models.transformers.fusion import BaseFuser, Fusers
+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,
diff --git a/vllm/model_executor/models/transformers/fuser.py b/vllm/model_executor/models/transformers/fuser.py
new file mode 100644
index 000000000000..0aa7c419ceb8
--- /dev/null
+++ b/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
diff --git a/vllm/model_executor/models/transformers/fusers/__init__.py b/vllm/model_executor/models/transformers/fusers/__init__.py
new file mode 100644
index 000000000000..4fad6befce62
--- /dev/null
+++ b/vllm/model_executor/models/transformers/fusers/__init__.py
@@ -0,0 +1,10 @@
+# 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.qkv import QKVFuser
+from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser
+
+__all__ = ["BaseFuser", "StackedFuser", "GLUFuser", "QKVFuser", "RMSNormFuser"]
diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py
new file mode 100644
index 000000000000..458f559018dc
--- /dev/null
+++ b/vllm/model_executor/models/transformers/fusers/base.py
@@ -0,0 +1,130 @@
+# 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.
+    """
+
+    @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`."""
+
+    fused_forward: Callable = field(init=False, repr=False)
+    """The compiled rewritten forward, set by `update_forward`."""
+
+    @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
diff --git a/vllm/model_executor/models/transformers/fusers/glu.py b/vllm/model_executor/models/transformers/fusers/glu.py
new file mode 100644
index 000000000000..52c985368103
--- /dev/null
+++ b/vllm/model_executor/models/transformers/fusers/glu.py
@@ -0,0 +1,193 @@
+# 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,
+    is_linear,
+    recover_forward,
+    replace_expr,
+    single_self_call,
+)
+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
+    merged_name: ClassVar[str] = "gate_up_proj"
+
+    @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] | None:
+        """Search graph for the GLU pattern `act(gate(x)) * up(x)`."""
+        for node in graph.nodes:
+            if (
+                node.op == "call_function"
+                and node.target == operator.mul
+                and len(node.args) == 2
+                and all(isinstance(arg, fx.Node) for arg in node.args)
+            ):
+                a, b = node.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
+        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 = 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
+        ):
+            return cls(act_node.target, gate_node.target, up_node.target)
+        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)
diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py
new file mode 100644
index 000000000000..743218b59b69
--- /dev/null
+++ b/vllm/model_executor/models/transformers/fusers/qkv.py
@@ -0,0 +1,195 @@
+# 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.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
+    merged_name: ClassVar[str] = "qkv_proj"
+
+    @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
+        return cls(q_name=q.target, k_name=k.target, v_name=v.target)
+
+    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)
diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py
new file mode 100644
index 000000000000..6ff5e0b069aa
--- /dev/null
+++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py
@@ -0,0 +1,139 @@
+# 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."""
+
+import operator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+from torch import fx, nn
+
+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 is_fn, is_method, peel
+
+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_method(node, "pow") or is_fn(node, torch.pow) or is_fn(node, operator.pow):
+        base, exp = node.args
+        return peel(base) is x and exp == 2
+    if is_method(node, "square") or is_fn(node, torch.square):
+        return peel(node.args[0]) is x
+    if is_fn(node, operator.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_fn(add, operator.add) or is_method(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_method(mean, "mean") or is_fn(mean, torch.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_fn(node, operator.add) or is_method(node, "add")):
+        return False
+    return any(isinstance(a, (int, float)) and a == 1 for a in node.args)
+
+
+@dataclass
+class RMSNormFuser(BaseFuser):
+    """Fuses an HF RMSNorm's `x * rsqrt(mean(x**2) + eps)` normalize into vLLM's
+    fused RMSNorm CustomOp, read structurally from the module's dataflow.
+
+    eps and the weight form (plain, zero-centered or absent) come straight from
+    the graph, so detection does not depend on the class name.
+    """
+
+    eps: float
+    has_weight: bool
+    zero_centered: bool
+    """Gemma-style `(1 + weight)` scaling (weight initialised at zero)."""
+
+    @classmethod
+    def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
+        x = next((n for n in graph.nodes if n.op == "placeholder"), None)
+        if x is None:
+            return None
+        # The rsqrt over the mean-square variance is the spine of the norm.
+        eps = rsqrt = None
+        for node in graph.nodes:
+            if (is_fn(node, torch.rsqrt) or is_method(node, "rsqrt")) and (
+                eps := _variance_eps(node, x)
+            ) is not None:
+                rsqrt = node
+                break
+        if rsqrt is None:
+            return None
+        # The `x * rsqrt(...)` normalize multiply.
+        normalize = next(
+            (
+                node
+                for node in graph.nodes
+                if is_fn(node, operator.mul) and rsqrt in map(peel, node.args)
+            ),
+            None,
+        )
+        if normalize is None:
+            return None
+        # An optional later `weight * normalized` (or `(1 + weight) * normalized`).
+        has_weight = zero_centered = False
+        for node in graph.nodes:
+            if not is_fn(node, operator.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)
+                has_weight = True
+                zero_centered = _is_one_plus(weight)
+                break
+        return cls(eps=eps, has_weight=has_weight, zero_centered=zero_centered)
+
+    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:
+        """Build vLLM's RMSNorm from the detected parameters.
+
+        The scale keeps its checkpoint name (`weight`), so no weight remapping is
+        needed. `var_hidden_size` has no Transformers equivalent."""
+        weight = getattr(module, "weight", None)
+        hidden_size = (
+            weight.size(0) if weight is not None else model_config.get_hidden_size()
+        )
+        if self.zero_centered:
+            return GemmaRMSNorm(hidden_size=hidden_size, eps=self.eps)
+        has_weight = self.has_weight and weight is not None
+        return RMSNorm(
+            hidden_size=hidden_size,
+            eps=self.eps,
+            has_weight=has_weight,
+            dtype=weight.dtype if has_weight else None,
+        )
diff --git a/vllm/model_executor/models/transformers/fusion.py b/vllm/model_executor/models/transformers/fusion.py
deleted file mode 100644
index 38c7b2da4c4c..000000000000
--- a/vllm/model_executor/models/transformers/fusion.py
+++ /dev/null
@@ -1,667 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Fusion patterns for the Transformers modeling backend.
-
-Each module class is traced once with `torch.fx` to detect a pattern from its
-dataflow, then its forward *source* is rewritten (AST) to replace only the
-matched calls; the rest stays live Python. The rewritten forward is compiled
-once per class and bound to each fused instance in place.
-"""
-
-import ast
-import inspect
-import operator
-import textwrap
-import types
-from abc import ABC, abstractmethod
-from collections import UserDict
-from collections.abc import Callable
-from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, ClassVar
-
-from cachetools import cached
-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,
-    QKVParallelLinear,
-)
-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())
-
-
-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 _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 _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 _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
-    )
-
-
-@dataclass
-class BaseFuser(ABC):
-    """A per-class fusion plan and the steps to apply it.
-
-    `match` and `update_forward` analyse the module *class* once (cached, see
-    `get_fuser`); `fuse` then 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`."""
-
-    fused_forward: Callable = field(init=False, repr=False)
-    """The compiled rewritten forward, set by `update_forward`."""
-
-    @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]}
-
-    @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 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 validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
-        """Whether this fuser can be applied to this `module` instance."""
-
-    @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
-
-
-@dataclass
-class GLUFuser(BaseFuser):
-    """Fuser for the GLU pattern `act(gate(x)) * up(x)`."""
-
-    act_name: str
-    gate_name: str
-    up_name: str
-    merged_name: ClassVar[str] = "gate_up_proj"
-
-    @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] | None:
-        """Search graph for the GLU pattern `act(gate(x)) * up(x)`."""
-        for node in graph.nodes:
-            if (
-                node.op == "call_function"
-                and node.target == operator.mul
-                and len(node.args) == 2
-                and all(isinstance(arg, fx.Node) for arg in node.args)
-            ):
-                a, b = node.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
-        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 = 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
-        ):
-            return cls(act_node.target, gate_node.target, up_node.target)
-        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)
-
-
-@dataclass
-class QKVFuser(BaseFuser):
-    """Fuser for the attention QKV pattern `q(x), k(x), v(x)`."""
-
-    q_name: str
-    k_name: str
-    v_name: str
-    merged_name: ClassVar[str] = "qkv_proj"
-
-    @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
-        return cls(q_name=q.target, k_name=k.target, v_name=v.target)
-
-    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)
-
-
-@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."""
-    # Currently, all patterns require at least 2 linears, this is a cheap early skip
-    if sum(isinstance(c, nn.Linear) for c in module.children()) < 2:
-        return None
-    tracer = _AllLeafTracer()
-    try:
-        graph = tracer.trace(module)
-    except Exception as exc:
-        # The graph is only evidence for matching, and the patterns appear at
-        # the top of their forwards, so the partial graph traced before the
-        # failure can still be matched.
-        logger.debug("Could not fully trace %s for fusion: %s", type(module), exc)
-        if (graph := getattr(tracer, "graph", None)) is None:
-            return None
-    for fuser_cls in (GLUFuser, QKVFuser):
-        if (fuser := fuser_cls.match(graph, module)) is not None:
-            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
-    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
diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py
new file mode 100644
index 000000000000..ed2decd850b5
--- /dev/null
+++ b/vllm/model_executor/models/transformers/fx_utils.py
@@ -0,0 +1,239 @@
+# 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
+
+from torch import fx, nn
+
+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 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 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:
+    return (
+        isinstance(node, fx.Node)
+        and node.op == "call_function"
+        and node.target is target
+    )
+
+
+def is_method(node: object, name: str) -> bool:
+    return (
+        isinstance(node, fx.Node) and node.op == "call_method" and node.target == name
+    )

From f12eb72ae3c9535f1e0c0ebfb0516432f0920c69 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:47:59 +0000
Subject: [PATCH 09/35] Use Fuser for RMSNorm too

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/base.py               |  7 +---
 .../models/transformers/utils.py              | 41 -------------------
 2 files changed, 1 insertion(+), 47 deletions(-)

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 5eac3eed3320..6c4220a210d0 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -59,7 +59,6 @@
     log_replacement,
     replace_conv_class,
     replace_linear_class,
-    replace_rms_norm_class,
 )
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
@@ -434,7 +433,7 @@ def recursive_replace(self):
         - Attention QKV projections with a fused `QKVParallelLinear` + split
         - `nn.Linear` with vLLM's tensor parallel linear classes
         - `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d`
-        - `*RMSNorm` with vLLM's `RMSNorm`
+        - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm`
         """
         tp_plan = self.model.tp_plan
 
@@ -498,10 +497,6 @@ def _recursive_replace(module: nn.Module, prefix: str):
                     )
                 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(
diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py
index 0a4ca94c5e97..4d9b01ce3930 100644
--- a/vllm/model_executor/models/transformers/utils.py
+++ b/vllm/model_executor/models/transformers/utils.py
@@ -23,10 +23,8 @@
 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",

From 16ff001cd4136f7ccd35bcdd95079af4d0ca367d Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:48:08 +0000
Subject: [PATCH 10/35] update tests

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/test_transformers_fusion.py | 151 ++++++++++++++++++++---
 1 file changed, 137 insertions(+), 14 deletions(-)

diff --git a/tests/models/test_transformers_fusion.py b/tests/models/test_transformers_fusion.py
index 76e93489e094..514089536491 100644
--- a/tests/models/test_transformers_fusion.py
+++ b/tests/models/test_transformers_fusion.py
@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the Transformers backend's fx-based fusion (GLU and QKV).
+"""Unit tests for the Transformers backend's fx-based fusers (GLU, QKV, RMSNorm).
 
 These exercise the structural detection and surgical subgraph rewrite in
-`vllm.model_executor.models.transformers.fusion` without needing a distributed
+`vllm.model_executor.models.transformers.fuser` without needing a distributed
 environment; the real `MergedColumnParallelLinear` / `QKVParallelLinear` /
 `...AndMul` integration is covered end-to-end by `test_fusion`, `test_models`
 and `test_quantization` in `tests/models/test_transformers.py`.
@@ -20,10 +20,10 @@
 
 # Registers the "vllm" attention interface in ALL_ATTENTION_FUNCTIONS
 import vllm.model_executor.models.transformers.base  # noqa: F401
-from vllm.model_executor.models.transformers import fusion
-from vllm.model_executor.models.transformers.fusion import (
+from vllm.model_executor.models.transformers.fuser import (
     GLUFuser,
     QKVFuser,
+    RMSNormFuser,
     get_fuser,
 )
 
@@ -186,9 +186,9 @@ def forward(self, q, k, v):
 
 @pytest.fixture(autouse=True)
 def _clear_fuser_cache():
-    fusion.get_fuser.cache_clear()
+    get_fuser.cache_clear()
     yield
-    fusion.get_fuser.cache_clear()
+    get_fuser.cache_clear()
 
 
 def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser):
@@ -316,7 +316,7 @@ def test_fuser_is_cached_per_class():
         fuser_a = get_fuser(GLUMLP())
         fuser_b = get_fuser(GLUMLP())
     assert fuser_a is fuser_b
-    assert GLUMLP in fusion.get_fuser.cache
+    assert GLUMLP in get_fuser.cache
 
 
 @pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP])
@@ -412,14 +412,137 @@ def test_act_and_mul_derived_from_module(default_vllm_config):
 
     from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul
 
-    assert isinstance(fusion.GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul)
-    assert isinstance(fusion.GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul)
-    gelu_tanh = fusion.GLUFuser._get_act_and_mul(GELUTanh())
+    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 = fusion.GLUFuser._get_act_and_mul(nn.GELU())
+    gelu = GLUFuser._get_act_and_mul(nn.GELU())
     assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none"
     # Not activations at all -> no fusion
-    assert fusion.GLUFuser._get_act_and_mul_name(nn.Dropout()) is None
-    assert fusion.GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None
+    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"):
-        fusion.GLUFuser._get_act_and_mul(nn.Dropout())
+        GLUFuser._get_act_and_mul(nn.Dropout())
+
+
+class LlamaRMSNorm(nn.Module):
+    """The canonical HF RMSNorm: `weight * x * rsqrt(mean(x**2) + eps)`."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class GemmaRMSNorm(nn.Module):
+    """Zero-centered weight: `(1 + weight) * normalized`."""
+
+    def __init__(self, dim: int = 16, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float())
+        output = output * (1.0 + self.weight.float())
+        return output.type_as(x)
+
+
+class WeightlessRMSNorm(nn.Module):
+    """No scale parameter (e.g. Gemma3n `with_scale=False`)."""
+
+    def __init__(self, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, x):
+        xf = x.float()
+        normed = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
+        return normed.type_as(x)
+
+
+class T5LayerNorm(nn.Module):
+    """An RMSNorm whose class name is *not* `*RMSNorm` (name-independence)."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states
+
+
+class NotAnRMSNorm(nn.Module):
+    """Mean-subtracting LayerNorm-like math -> not an RMSNorm."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden))
+        self.eps = eps
+
+    def forward(self, x):
+        x = x - x.mean(-1, keepdim=True)
+        return self.weight * x / torch.sqrt(x.var(-1, keepdim=True) + self.eps)
+
+
+@pytest.mark.parametrize(
+    "cls,eps,has_weight,zero_centered",
+    [
+        (LlamaRMSNorm, 1e-5, True, False),
+        (GemmaRMSNorm, 1e-6, True, True),
+        (WeightlessRMSNorm, 1e-6, False, False),
+        (T5LayerNorm, 1e-6, True, False),
+    ],
+)
+def test_detects_rms_norm_variants(cls, eps, has_weight, zero_centered):
+    with torch.device("meta"):
+        fuser = get_fuser(cls())
+    assert isinstance(fuser, RMSNormFuser)
+    # eps and the weight form are read from the graph, not the class name.
+    assert fuser.eps == eps
+    assert fuser.has_weight == has_weight
+    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,expected,zero_centered",
+    [
+        (LlamaRMSNorm, "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)
+    types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm}
+    assert type(built) is types_by_name[expected]
+    assert built.variance_epsilon == fuser.eps

From 6c4debd50f3729345ff4fcc3f4a33958f22b2f1e Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Wed, 1 Jul 2026 16:21:06 +0000
Subject: [PATCH 11/35] Log fusion info

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 vllm/model_executor/models/transformers/base.py    |  1 +
 .../models/transformers/fusers/base.py             | 12 ++++++++++++
 .../models/transformers/fusers/glu.py              |  1 +
 .../models/transformers/fusers/qkv.py              |  1 +
 .../models/transformers/fusers/rms_norm.py         | 14 +++++++++++++-
 5 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 6c4220a210d0..c0864dd4c319 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -502,6 +502,7 @@ def _recursive_replace(module: nn.Module, prefix: str):
                     new_module = fuser.fuse(
                         child_module, qual_name, self.model_config, self.quant_config
                     )
+                    logger.info_once(fuser.info)
                     _recursive_replace(new_module, prefix=qual_name)
                 else:
                     _recursive_replace(child_module, prefix=qual_name)
diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py
index 458f559018dc..b39075c3785b 100644
--- a/vllm/model_executor/models/transformers/fusers/base.py
+++ b/vllm/model_executor/models/transformers/fusers/base.py
@@ -26,6 +26,11 @@ class BaseFuser(ABC):
     module to install in its place.
     """
 
+    @property
+    @abstractmethod
+    def info(self) -> str:
+        """A human-readable description of the fusion, for logging."""
+
     @classmethod
     @abstractmethod
     def match(cls, graph: fx.Graph, module: nn.Module) -> "BaseFuser | None":
@@ -70,10 +75,17 @@ class StackedFuser(BaseFuser):
 
     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)."""
 
     fused_forward: Callable = field(init=False, repr=False)
     """The compiled rewritten forward, set by `update_forward`."""
 
+    @property
+    def info(self) -> str:
+        sources = " + ".join(name for name, _ in self.shards)
+        return f"Fused {sources} -> {self.merged_name} ({self.merged_cls})"
+
     @property
     @abstractmethod
     def shards(self) -> list[tuple[str, ShardId]]:
diff --git a/vllm/model_executor/models/transformers/fusers/glu.py b/vllm/model_executor/models/transformers/fusers/glu.py
index 52c985368103..00bebf25a7cd 100644
--- a/vllm/model_executor/models/transformers/fusers/glu.py
+++ b/vllm/model_executor/models/transformers/fusers/glu.py
@@ -50,6 +50,7 @@ class GLUFuser(StackedFuser):
     gate_name: str
     up_name: str
     merged_name: ClassVar[str] = "gate_up_proj"
+    merged_cls: ClassVar[str] = "MergedColumnParallelLinear"
 
     @property
     def shards(self) -> list[tuple[str, ShardId]]:
diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py
index 743218b59b69..70532e742283 100644
--- a/vllm/model_executor/models/transformers/fusers/qkv.py
+++ b/vllm/model_executor/models/transformers/fusers/qkv.py
@@ -36,6 +36,7 @@ class QKVFuser(StackedFuser):
     k_name: str
     v_name: str
     merged_name: ClassVar[str] = "qkv_proj"
+    merged_cls: ClassVar[str] = "QKVParallelLinear"
 
     @property
     def shards(self) -> list[tuple[str, ShardId]]:
diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py
index 6ff5e0b069aa..691fda446fc2 100644
--- a/vllm/model_executor/models/transformers/fusers/rms_norm.py
+++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py
@@ -70,6 +70,13 @@ class RMSNormFuser(BaseFuser):
     has_weight: bool
     zero_centered: bool
     """Gemma-style `(1 + weight)` scaling (weight initialised at zero)."""
+    source_cls: str
+    """Class name of the HF norm this was matched from (for logging)."""
+
+    @property
+    def info(self) -> str:
+        norm = "GemmaRMSNorm" if self.zero_centered else "RMSNorm"
+        return f"Fused {self.source_cls} -> {norm} (CustomOp)"
 
     @classmethod
     def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
@@ -108,7 +115,12 @@ def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
                 has_weight = True
                 zero_centered = _is_one_plus(weight)
                 break
-        return cls(eps=eps, has_weight=has_weight, zero_centered=zero_centered)
+        return cls(
+            eps=eps,
+            has_weight=has_weight,
+            zero_centered=zero_centered,
+            source_cls=type(module).__name__,
+        )
 
     def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
         return True

From 7d3a7514401e3e3cc4d82bf1f7a49030774622ae Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Wed, 1 Jul 2026 19:52:05 +0000
Subject: [PATCH 12/35] Use fusion mechanism for MoE and RMSNorm

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .buildkite/test_areas/models_basic.yaml       |   4 +-
 ..._fusion.py => test_transformers_fusers.py} | 268 +++++++++++++++++-
 .../models/transformers/base.py               |   2 +-
 .../models/transformers/fusers/base.py        |  18 +-
 .../models/transformers/fusers/glu.py         |   7 +-
 .../models/transformers/fusers/moe.py         | 268 ++++++++++++++++++
 .../models/transformers/fusers/qkv.py         |   7 +-
 .../models/transformers/fusers/rms_norm.py    |  38 +--
 .../models/transformers/fx_utils.py           |  18 ++
 .../model_executor/models/transformers/moe.py | 111 +++++---
 10 files changed, 662 insertions(+), 79 deletions(-)
 rename tests/models/{test_transformers_fusion.py => test_transformers_fusers.py} (68%)
 create mode 100644 vllm/model_executor/models/transformers/fusers/moe.py

diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml
index 8491987f964d..52b3330154c0 100644
--- a/.buildkite/test_areas/models_basic.yaml
+++ b/.buildkite/test_areas/models_basic.yaml
@@ -60,7 +60,7 @@ steps:
   - vllm/
   - tests/models/test_utils.py
   - tests/models/test_vision.py
-  - tests/models/test_transformers_fusion.py
+  - tests/models/test_transformers_fusers.py
   device: cpu-small
   commands:
-    - pytest -v -s models/test_utils.py models/test_vision.py models/test_transformers_fusion.py
+    - pytest -v -s models/test_utils.py models/test_vision.py models/test_transformers_fusers.py
diff --git a/tests/models/test_transformers_fusion.py b/tests/models/test_transformers_fusers.py
similarity index 68%
rename from tests/models/test_transformers_fusion.py
rename to tests/models/test_transformers_fusers.py
index 514089536491..1e162bb8b804 100644
--- a/tests/models/test_transformers_fusion.py
+++ b/tests/models/test_transformers_fusers.py
@@ -1,13 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the Transformers backend's fx-based fusers (GLU, QKV, RMSNorm).
-
-These exercise the structural detection and surgical subgraph rewrite in
-`vllm.model_executor.models.transformers.fuser` without needing a distributed
-environment; the real `MergedColumnParallelLinear` / `QKVParallelLinear` /
-`...AndMul` integration is covered end-to-end by `test_fusion`, `test_models`
-and `test_quantization` in `tests/models/test_transformers.py`.
-"""
+"""Unit tests for the Transformers backend's fx fusers (MoE, GLU, QKV, RMSNorm)."""
 
 import inspect
 import types
@@ -26,6 +19,7 @@
     RMSNormFuser,
     get_fuser,
 )
+from vllm.model_executor.models.transformers.fusers.moe import MoEFuser
 
 
 class SiluAndMulStub(nn.Module):
@@ -546,3 +540,261 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c
     types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm}
     assert type(built) is types_by_name[expected]
     assert built.variance_epsilon == fuser.eps
+
+
+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/noaux router whose score-correction bias is a buffer (as in
+    DeepSeek-V3).
+
+    It is `weight`-only in its parameters, so it is declined via the correction
+    bias read structurally from the graph, not by the parameter check."""
+
+    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 (an extra `bias` parameter).
+
+    `build_gate` rebuilds the router as a bias-free `ReplicatedLinear`, so a
+    router carrying any other parameter cannot be reproduced faithfully."""
+
+    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 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):
+    """A single-tensor-returning MoE block (Qwen3-style)."""
+
+    def __init__(self, router_cls=TopKRouter):
+        super().__init__()
+        self.experts = MoEExperts()
+        self.gate = router_cls()
+
+    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).reshape(hidden_states.shape)
+
+
+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 forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        out = self.experts(x, index, weights)
+        out = out + torch.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
+        return out.reshape(hidden_states.shape)
+
+
+class MoEBlockSharedNoGate(MoEBlock):
+    """A block with an ungated shared expert -> native, shared passed through."""
+
+    def __init__(self):
+        super().__init__()
+        self.shared_expert = GLUMLP()
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        out = self.experts(x, index, weights) + self.shared_expert(x)
+        return out.reshape(hidden_states.shape)
+
+
+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 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 block whose shared expert is a non-GLU MLP -> detected by dataflow.
+
+    It is added to the experts' output, so it is recognised as the shared
+    expert even though it is not a GLU (no gate/up merge is needed).
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.shared_expert = PlainMLP()
+
+    def forward(self, hidden_states):
+        x = hidden_states.reshape(-1, hidden_states.shape[-1])
+        _, weights, index = self.gate(x)
+        out = self.experts(x, index, weights) + self.shared_expert(x)
+        return out.reshape(hidden_states.shape)
+
+
+class MoEBlockUnaccounted(MoEBlock):
+    """A block with a weight-bearing child outside the experts + shared-expert
+    dataflow (here a pre-router transform) -> must decline, since the rewritten
+    forward would drop it."""
+
+    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(MoEBlock):
+    """Like `MoEBlockUnaccounted`, but the dropped child holds only a buffer, so
+    the fail-closed check must inspect buffers as well as parameters."""
+
+    def __init__(self):
+        super().__init__()
+        self.extra = BufferScale()
+
+    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)
+
+
+@pytest.mark.parametrize("sigmoid", [False, True])
+def test_moe_fuser_detects_router(sigmoid):
+    with torch.device("meta"):
+        block = MoEBlock(lambda: TopKRouter(sigmoid=sigmoid))
+    fuser = MoEFuser.match(block)
+    assert isinstance(fuser, MoEFuser)
+    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 = MoEFuser.match(block)
+    assert isinstance(fuser, MoEFuser)
+    assert fuser.shared_name == "shared_expert"
+    assert fuser.shared_gate_name == "shared_expert_gate"
+    # The shared expert's gate/up projections are merged, scoped to its qualname.
+    assert fuser.orig_to_new_stacked("model.layers.0.mlp") == {
+        "model.layers.0.mlp.shared_expert.gate_proj": (
+            "model.layers.0.mlp.shared_expert.gate_up_proj",
+            0,
+        ),
+        "model.layers.0.mlp.shared_expert.up_proj": (
+            "model.layers.0.mlp.shared_expert.gate_up_proj",
+            1,
+        ),
+    }
+
+
+def test_moe_fuser_shared_without_gate():
+    with torch.device("meta"):
+        block = MoEBlockSharedNoGate()
+    fuser = MoEFuser.match(block)
+    assert isinstance(fuser, MoEFuser)
+    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 = MoEFuser.match(block)
+    assert isinstance(fuser, MoEFuser)
+    # 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
+    # A non-GLU shared expert needs no gate/up merge; it loads by identity.
+    assert fuser.shared_glu is None
+    assert fuser.orig_to_new_stacked("model.layers.0.mlp") == {}
+
+
+@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)
+        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 MoEFuser.match(block) is None
diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index c0864dd4c319..c55cdf871db6 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -502,7 +502,7 @@ def _recursive_replace(module: nn.Module, prefix: str):
                     new_module = fuser.fuse(
                         child_module, qual_name, self.model_config, self.quant_config
                     )
-                    logger.info_once(fuser.info)
+                    logger.info_once(fuser.info(child_name))
                     _recursive_replace(new_module, prefix=qual_name)
                 else:
                     _recursive_replace(child_module, prefix=qual_name)
diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py
index b39075c3785b..54fb2d09165c 100644
--- a/vllm/model_executor/models/transformers/fusers/base.py
+++ b/vllm/model_executor/models/transformers/fusers/base.py
@@ -26,10 +26,9 @@ class BaseFuser(ABC):
     module to install in its place.
     """
 
-    @property
     @abstractmethod
-    def info(self) -> str:
-        """A human-readable description of the fusion, for logging."""
+    def info(self, name: str) -> str:
+        """A human-readable description of the fusion at `name`, for logging."""
 
     @classmethod
     @abstractmethod
@@ -78,13 +77,18 @@ class StackedFuser(BaseFuser):
     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`."""
 
-    @property
-    def info(self) -> str:
-        sources = " + ".join(name for name, _ in self.shards)
-        return f"Fused {sources} -> {self.merged_name} ({self.merged_cls})"
+    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
diff --git a/vllm/model_executor/models/transformers/fusers/glu.py b/vllm/model_executor/models/transformers/fusers/glu.py
index 00bebf25a7cd..1536b1df5ff1 100644
--- a/vllm/model_executor/models/transformers/fusers/glu.py
+++ b/vllm/model_executor/models/transformers/fusers/glu.py
@@ -124,7 +124,12 @@ def match(cls, graph: fx.Graph, module: nn.Module) -> "GLUFuser | None":
         if gate.in_features == up.in_features and (gate.bias is None) == (
             up.bias is None
         ):
-            return cls(act_node.target, gate_node.target, up_node.target)
+            return cls(
+                source_cls=type(module).__name__,
+                act_name=act_node.target,
+                gate_name=gate_node.target,
+                up_name=up_node.target,
+            )
         return None
 
     def update_forward(self, module: nn.Module) -> None:
diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py
new file mode 100644
index 000000000000..15a58b69bc37
--- /dev/null
+++ b/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.fusers.glu import GLUFuser
+from vllm.model_executor.models.transformers.fx_utils import (
+    find_node,
+    is_op,
+    peel,
+    trace,
+)
+from vllm.model_executor.models.utils import (
+    ShardId,
+    maybe_prefix,
+    sequence_parallel_chunk,
+)
+
+
+def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]:
+    """`module`'s own state -- its named parameters and buffers together."""
+    return chain(module.named_parameters(), module.named_buffers())
+
+
+def _returns_tuple(cls: type[nn.Module]) -> bool:
+    """True if the class's forward returns a tuple (not a single tensor)."""
+    try:
+        source = textwrap.dedent(inspect.getsource(inspect.unwrap(cls.forward)))
+        tree = ast.parse(source)
+    except (OSError, SyntaxError, TypeError):
+        return True
+    return any(
+        isinstance(node, ast.Return) and isinstance(node.value, ast.Tuple)
+        for node in ast.walk(tree)
+    )
+
+
+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 _match_router(gate: nn.Module) -> str | None:
+    """Detects a router by dataflow: linear + softmax/sigmoid + top-k."""
+    if [name for name, _ in named_state(gate)] != ["weight"]:
+        return None
+    graph = trace(gate)
+    if graph is None:
+        return None
+    nodes = list(graph.nodes)
+    if not any(is_op(n, "linear") for n in nodes):
+        return None
+    if not any(is_op(n, "topk") for n in nodes):
+        return None
+    softmax = any(is_op(n, "softmax") for n in nodes)
+    sigmoid = any(is_op(n, "sigmoid") for n in nodes)
+    if softmax == sigmoid:  # need exactly one scoring function
+        return None
+    return "softmax" if softmax else "sigmoid"
+
+
+def _reaches(node: fx.Node, key: str) -> set[fx.Node]:
+    """`node` and every fx node connected to it via `key` (`users`/`inputs`)."""
+    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
+
+
+def _match_shared_expert(
+    graph: fx.Graph, experts: str = "experts"
+) -> 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 node in graph.nodes:
+        if not is_op(node, "add"):
+            continue
+        operands = [a for a in node.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
+
+
+class SharedExpertMLP(nn.Module):
+    """Wraps an HF shared expert, applying the output gating it is paired with."""
+
+    def __init__(self, shared_expert: nn.Module, gate: nn.Module | None = None):
+        super().__init__()
+        self.shared_expert = shared_expert
+        self.gate = gate
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        out = self.shared_expert(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:
+    """Native MoE block forward: the internal router routes + runs shared experts.
+
+    Under sequence-parallel MoE, tokens are scattered across TP ranks before the
+    experts and all-gathered back after (matching vLLM's native MoE blocks).
+    """
+    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 MoEFuser:
+    """Fuser for an HF MoE block, routing through vLLM's `FusedMoE` with the
+    same shared expert (if any) and gate (if any) as the original block."""
+
+    gate_name: str
+    scoring_func: str
+    shared_name: str | None
+    shared_gate_name: str | None
+    shared_glu: GLUFuser | None
+
+    @classmethod
+    def match(cls, mlp: nn.Module) -> "MoEFuser | None":
+        # The native block forward returns a single tensor.
+        if _returns_tuple(type(mlp)):
+            return None
+        # Router: the child that scores + top-k selects (traced per child).
+        gate_name = scoring_func = None
+        for name, child in mlp.named_children():
+            if name != "experts" and (func := _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: whatever the block adds to the experts' output, found by
+        # dataflow so any MLP works, not just a GLU. A GLU inside it is fused by
+        # the base pass when it descends into the block child; we only precompute
+        # its gate/up remap (see `orig_to_new_stacked`).
+        graph = trace(mlp)
+        if graph is None:
+            return None
+        shared_name, shared_gate_name = _match_shared_expert(graph)
+        if shared_gate_name is not None and not _is_scalar_gate(
+            getattr(mlp, shared_gate_name)
+        ):
+            return None
+        shared_glu = None
+        if (
+            shared_name is not None
+            and (sgraph := trace(shared := getattr(mlp, shared_name))) is not None
+        ):
+            shared_glu = GLUFuser.match(sgraph, shared)
+        # Fail closed: `rewrite_forward` runs only the experts and the detected
+        # shared expert, so any other stateful child would be dropped.
+        accounted = {"experts", gate_name, shared_name, shared_gate_name}
+        for name, child in mlp.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, shared_glu)
+
+    def build_gate(self, mlp: nn.Module, prefix: str) -> ReplicatedLinear:
+        """Rebuild the HF router as a `ReplicatedLinear` producing logits.
+
+        Shapes come from the original router weight (`[num_experts, hidden]`),
+        the plain `F.linear` weight it loads by identity into. Left unquantized
+        -- routers run in full precision -- matching the checkpoint's weight.
+        """
+        num_experts, hidden_size = getattr(mlp, self.gate_name).weight.shape
+        gate = ReplicatedLinear(
+            hidden_size,
+            num_experts,
+            bias=False,
+            prefix=maybe_prefix(prefix, self.gate_name),
+        )
+        setattr(mlp, self.gate_name, gate)
+        return gate
+
+    def build_shared_experts(
+        self, mlp: nn.Module, prefix: str
+    ) -> SharedExpertMLP | None:
+        """Wrap the HF shared expert, folding in its sigmoid output gate.
+
+        The MLP is the original module (converted in place by the base pass); the
+        sibling gate is rebuilt as a `ReplicatedLinear` (its `[1, hidden]` weight
+        is used as a plain `F.linear` weight) so the wrapper holds a reference
+        that survives the base pass's replacement, and it loads by identity.
+        """
+        if self.shared_name is None:
+            return None
+        shared_expert = getattr(mlp, self.shared_name)
+        gate = None
+        if self.shared_gate_name is not None:
+            hf_gate = getattr(mlp, 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(mlp, self.shared_gate_name, gate)
+        return SharedExpertMLP(shared_expert, gate)
+
+    def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
+        """Merge the shared expert's gate/up projections (scoped to its qualname).
+
+        Registered here so loading is correct regardless of whether the base pass
+        reaches the shared expert via the block child or the runner alias first.
+        """
+        if self.shared_name is None or self.shared_glu is None:
+            return {}
+        return self.shared_glu.orig_to_new_stacked(
+            maybe_prefix(prefix, self.shared_name)
+        )
+
+    def rewrite_forward(self, mlp: nn.Module) -> None:
+        """Bind the native block forward (`experts` routes + runs shared)."""
+        mlp.forward = types.MethodType(_moe_block_forward, mlp)
diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py
index 70532e742283..5de0bdca405a 100644
--- a/vllm/model_executor/models/transformers/fusers/qkv.py
+++ b/vllm/model_executor/models/transformers/fusers/qkv.py
@@ -78,7 +78,12 @@ 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
-        return cls(q_name=q.target, k_name=k.target, v_name=v.target)
+        return cls(
+            source_cls=type(module).__name__,
+            q_name=q.target,
+            k_name=k.target,
+            v_name=v.target,
+        )
 
     def update_forward(self, module: nn.Module) -> None:
         """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source."""
diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py
index 691fda446fc2..04797c574955 100644
--- a/vllm/model_executor/models/transformers/fusers/rms_norm.py
+++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py
@@ -2,16 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """RMSNorm fuser: detect the norm structurally and swap in vLLM's fused RMSNorm."""
 
-import operator
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-import torch
 from torch import fx, nn
 
 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 is_fn, is_method, peel
+from vllm.model_executor.models.transformers.fx_utils import find_node, is_op, peel
 
 if TYPE_CHECKING:
     from vllm.config.model import ModelConfig
@@ -21,12 +19,12 @@
 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_method(node, "pow") or is_fn(node, torch.pow) or is_fn(node, operator.pow):
+    if is_op(node, "pow"):
         base, exp = node.args
         return peel(base) is x and exp == 2
-    if is_method(node, "square") or is_fn(node, torch.square):
+    if is_op(node, "square"):
         return peel(node.args[0]) is x
-    if is_fn(node, operator.mul):
+    if is_op(node, "mul"):
         a, b = node.args
         return peel(a) is x and peel(b) is x
     return False
@@ -35,14 +33,14 @@ def _is_squared(node: object, x: fx.Node) -> bool:
 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_fn(add, operator.add) or is_method(add, "add")):
+    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_method(mean, "mean") or is_fn(mean, torch.mean)):
+    if not is_op(mean, "mean"):
         return None
     if not _is_squared(mean.args[0], x):
         return None
@@ -52,7 +50,7 @@ def _variance_eps(rsqrt: fx.Node, x: fx.Node) -> float | None:
 def _is_one_plus(node: object) -> bool:
     """`1 + weight` in either operand order (marks a zero-centered weight)."""
     node = peel(node)
-    if not (is_fn(node, operator.add) or is_method(node, "add")):
+    if not is_op(node, "add"):
         return False
     return any(isinstance(a, (int, float)) and a == 1 for a in node.args)
 
@@ -73,41 +71,33 @@ class RMSNormFuser(BaseFuser):
     source_cls: str
     """Class name of the HF norm this was matched from (for logging)."""
 
-    @property
-    def info(self) -> str:
+    def info(self, name: str) -> str:
         norm = "GemmaRMSNorm" if self.zero_centered else "RMSNorm"
-        return f"Fused {self.source_cls} -> {norm} (CustomOp)"
+        return f"Fused: {name} ({self.source_cls}) -> {norm} (CustomOp)"
 
     @classmethod
     def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
-        x = next((n for n in graph.nodes if n.op == "placeholder"), None)
+        x = find_node(graph, lambda n: n.op == "placeholder")
         if x is None:
             return None
         # The rsqrt over the mean-square variance is the spine of the norm.
         eps = rsqrt = None
         for node in graph.nodes:
-            if (is_fn(node, torch.rsqrt) or is_method(node, "rsqrt")) and (
-                eps := _variance_eps(node, x)
-            ) is not None:
+            if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None:
                 rsqrt = node
                 break
         if rsqrt is None:
             return None
         # The `x * rsqrt(...)` normalize multiply.
-        normalize = next(
-            (
-                node
-                for node in graph.nodes
-                if is_fn(node, operator.mul) and rsqrt in map(peel, node.args)
-            ),
-            None,
+        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`).
         has_weight = zero_centered = False
         for node in graph.nodes:
-            if not is_fn(node, operator.mul) or node is normalize:
+            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:
diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py
index ed2decd850b5..51310aa126b4 100644
--- a/vllm/model_executor/models/transformers/fx_utils.py
+++ b/vllm/model_executor/models/transformers/fx_utils.py
@@ -14,7 +14,9 @@
 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
 
@@ -204,6 +206,11 @@ def visit(self, node: ast.AST) -> ast.AST:
     _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(
@@ -226,6 +233,7 @@ def peel(node: object) -> object:
 
 
 def is_fn(node: object, target: Callable) -> bool:
+    """Is node `<target>()`."""
     return (
         isinstance(node, fx.Node)
         and node.op == "call_function"
@@ -234,6 +242,16 @@ def is_fn(node: object, target: Callable) -> bool:
 
 
 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))
diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py
index d5267a261797..a309aed3bd9c 100644
--- a/vllm/model_executor/models/transformers/moe.py
+++ b/vllm/model_executor/models/transformers/moe.py
@@ -26,9 +26,11 @@
 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 MoEFuser
 from vllm.model_executor.models.utils import maybe_prefix
 from vllm.utils.torch_utils import direct_register_custom_op
 
@@ -37,6 +39,8 @@
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
 
+logger = init_logger(__name__)
+
 
 @dataclass
 class TransformersMoEState:
@@ -221,6 +225,9 @@ def _recursive_replace(module: nn.Module, prefix: str):
                     # Alias for readability
                     mlp = module
                     experts = child_module
+                    # Class of the fused block (parent of gate/experts/shared)
+                    mlp_cls = type(mlp).__name__
+                    experts_cls = type(experts).__name__
                     # Do the experts have biases
                     has_bias = False
                     for experts_param_name, _ in experts.named_parameters():
@@ -235,55 +242,89 @@ def _recursive_replace(module: nn.Module, prefix: str):
                                 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},
                     )
+                    if self.num_expert_groups <= 1 and (fuser := MoEFuser.match(mlp)):
+                        # MLP forward is fully replaced.
+                        # gate/router and shared expert (if any) runs in FusedMoE.
+                        gate = fuser.build_gate(mlp, prefix)
+                        shared_experts = fuser.build_shared_experts(mlp, prefix)
+                        if shared_experts is not None:
+                            self.hf_to_vllm_mapper.orig_to_new_stacked.update(
+                                fuser.orig_to_new_stacked(prefix)
+                            )
+                        kwargs |= dict(
+                            scoring_func=fuser.scoring_func,
+                            is_sequence_parallel=(
+                                self.parallel_config.use_sequence_parallel_moe
+                            ),
+                            gate=gate,
+                            shared_experts=shared_experts,
+                        )
+                        fuser.rewrite_forward(mlp)
+                        routed = "gate + experts"
+                        if fuser.shared_name:
+                            routed += " + shared experts"
+                        logger.info_once(
+                            "Fused: %s (%s) -> FusedMoE (internal routing)",
+                            routed,
+                            mlp_cls,
+                        )
+                    else:
+                        # MLP 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)
                     mlp.experts = fused_experts
                     log_replacement(qual_name, experts, fused_experts)
                     # Update MixtureOfExperts mixin state

From 83da8be09526e98f17c9e62e4c07f9e304ef2481 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 10:10:16 +0000
Subject: [PATCH 13/35] Update perf expectations in doc

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 docs/models/supported_models.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 87edb27f1fab..215fd059f334 100644
--- a/docs/models/supported_models.md
+++ b/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:
 

From f9a601bfea1334a0df60039e10cd6dce20f9ccd1 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 10:10:31 +0000
Subject: [PATCH 14/35] Support torch.nn.RMSNorm

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/fusers/rms_norm.py    | 26 ++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py
index 04797c574955..1255ffc56a99 100644
--- a/vllm/model_executor/models/transformers/fusers/rms_norm.py
+++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py
@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+import torch
 from torch import fx, nn
 
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
@@ -64,7 +65,8 @@ class RMSNormFuser(BaseFuser):
     the graph, so detection does not depend on the class name.
     """
 
-    eps: float
+    eps: float | None
+    """`None` only for a fused `rms_norm` op with default eps; resolved in `fuse`."""
     has_weight: bool
     zero_centered: bool
     """Gemma-style `(1 + weight)` scaling (weight initialised at zero)."""
@@ -80,6 +82,19 @@ def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
         x = find_node(graph, lambda n: n.op == "placeholder")
         if x is None:
             return None
+        # Handle native torch `rms_norm` op.
+        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
+            weight = args[2] if len(args) > 2 else kwargs.get("weight")
+            eps = args[3] if len(args) > 3 else kwargs.get("eps")
+            return cls(
+                eps=eps if isinstance(eps, (int, float)) else None,
+                has_weight=isinstance(weight, fx.Node),
+                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.
         eps = rsqrt = None
         for node in graph.nodes:
@@ -130,12 +145,17 @@ def fuse(
         hidden_size = (
             weight.size(0) if weight is not None else model_config.get_hidden_size()
         )
+        eps = self.eps
+        if eps is None:
+            # Could be `None` for native torch `rms_norm`. 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 GemmaRMSNorm(hidden_size=hidden_size, eps=self.eps)
+            return GemmaRMSNorm(hidden_size=hidden_size, eps=eps)
         has_weight = self.has_weight and weight is not None
         return RMSNorm(
             hidden_size=hidden_size,
-            eps=self.eps,
+            eps=eps,
             has_weight=has_weight,
             dtype=weight.dtype if has_weight else None,
         )

From 0cc28039b9a27e056598911883ea933d27271cce Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 10:16:19 +0000
Subject: [PATCH 15/35] Update tests

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/test_transformers_fusers.py | 117 ++++++++++++-----------
 1 file changed, 63 insertions(+), 54 deletions(-)

diff --git a/tests/models/test_transformers_fusers.py b/tests/models/test_transformers_fusers.py
index 1e162bb8b804..14c02498ca30 100644
--- a/tests/models/test_transformers_fusers.py
+++ b/tests/models/test_transformers_fusers.py
@@ -21,6 +21,8 @@
 )
 from vllm.model_executor.models.transformers.fusers.moe import MoEFuser
 
+# GLU and QKV fusers (shared module fixtures and fuser-application helpers)
+
 
 class SiluAndMulStub(nn.Module):
     """Stand-in for vLLM's `SiluAndMul` (no vLLM config required)."""
@@ -72,8 +74,7 @@ def __init__(self):
 
 
 class UntraceableMLP(GLUMLP):
-    """Data-dependent control flow *before* the pattern -> no match (the
-    partial trace ends before reaching the GLU)."""
+    """Data-dependent control flow *before* the GLU -> no match."""
 
     def forward(self, x):
         if x.sum() > 0:  # noqa: SIM108 - intentionally untraceable
@@ -82,8 +83,7 @@ def forward(self, x):
 
 
 class UntraceableTailGLUMLP(GLUMLP):
-    """Data-dependent control flow *after* the pattern -> still fusable (the
-    partial trace contains the GLU, and the AST rewrite keeps the tail live)."""
+    """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))
@@ -95,11 +95,9 @@ def forward(self, x):
 class FakeAttention(nn.Module):
     """Mimics the structure of a HF v5 attention module under vLLM.
 
-    The forward exercises the patterns the fusion tracer must support: shape
-    unpacking, the dead KV cache branch, and the `attention_interface` call
-    splatting `**kwargs` (resolved to vLLM's interface, which is traced
-    through).
-    """
+    The forward exercises the patterns the fusion tracer must support: shape unpacking,
+    the dead KV cache branch, and the `attention_interface` call splatting `**kwargs`
+    (resolved to vLLM's interface, which is traced through)."""
 
     is_causal = True
 
@@ -419,91 +417,82 @@ def test_act_and_mul_derived_from_module(default_vllm_config):
         GLUFuser._get_act_and_mul(nn.Dropout())
 
 
-class LlamaRMSNorm(nn.Module):
+# RMSNorm fuser
+
+
+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):
+    def __init__(self, hidden: int = 16, eps: float = 1e-5, weight: bool = True):
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden))
+        if weight:
+            self.weight = nn.Parameter(torch.ones(hidden))
         self.variance_epsilon = eps
 
-    def forward(self, hidden_states):
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
+    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(nn.Module):
-    """Zero-centered weight: `(1 + weight) * normalized`."""
 
-    def __init__(self, dim: int = 16, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.zeros(dim))
+class GemmaRMSNorm(RMSNorm):
+    """Zero-centered weight: `(1 + weight) * normalized`."""
 
-    def _norm(self, x):
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+    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):
-        output = self._norm(x.float())
-        output = output * (1.0 + self.weight.float())
-        return output.type_as(x)
+        return (1.0 + self.weight) * self._rms(x.to(torch.float32)).to(x.dtype)
 
 
-class WeightlessRMSNorm(nn.Module):
+class WeightlessRMSNorm(RMSNorm):
     """No scale parameter (e.g. Gemma3n `with_scale=False`)."""
 
-    def __init__(self, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__(hidden, eps, weight=False)
 
     def forward(self, x):
-        xf = x.float()
-        normed = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
-        return normed.type_as(x)
+        return self._rms(x.to(torch.float32)).to(x.dtype)
 
 
-class T5LayerNorm(nn.Module):
-    """An RMSNorm whose class name is *not* `*RMSNorm` (name-independence)."""
+class LayerNorm(RMSNorm):
+    """An RMSNorm whose class name is *not* `*RMSNorm` (name-independence), and
+    which keeps the input dtype throughout (no float32 upcast)."""
 
     def __init__(self, hidden: int = 16, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden))
-        self.variance_epsilon = eps
+        super().__init__(hidden, eps)
 
-    def forward(self, hidden_states):
-        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states
+    def forward(self, x):
+        return self.weight * self._rms(x)
 
 
-class NotAnRMSNorm(nn.Module):
+class NotAnRMSNorm(RMSNorm):
     """Mean-subtracting LayerNorm-like math -> not an RMSNorm."""
 
     def __init__(self, hidden: int = 16, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden))
-        self.eps = eps
+        super().__init__(hidden, eps)
 
     def forward(self, x):
         x = x - x.mean(-1, keepdim=True)
-        return self.weight * x / torch.sqrt(x.var(-1, keepdim=True) + self.eps)
+        variance = x.var(-1, keepdim=True)
+        return self.weight * x / torch.sqrt(variance + self.variance_epsilon)
 
 
 @pytest.mark.parametrize(
     "cls,eps,has_weight,zero_centered",
     [
-        (LlamaRMSNorm, 1e-5, True, False),
+        (RMSNorm, 1e-5, True, False),
         (GemmaRMSNorm, 1e-6, True, True),
         (WeightlessRMSNorm, 1e-6, False, False),
-        (T5LayerNorm, 1e-6, True, False),
+        (LayerNorm, 1e-6, True, False),
+        (torch.nn.RMSNorm, 1e-5, True, False),  # fused `F.rms_norm` op
     ],
 )
 def test_detects_rms_norm_variants(cls, eps, has_weight, zero_centered):
     with torch.device("meta"):
-        fuser = get_fuser(cls())
+        fuser = get_fuser(cls(16, eps=eps))
     assert isinstance(fuser, RMSNormFuser)
     # eps and the weight form are read from the graph, not the class name.
     assert fuser.eps == eps
@@ -521,7 +510,7 @@ def test_non_rms_norms_are_not_matched(cls):
 @pytest.mark.parametrize(
     "cls,expected,zero_centered",
     [
-        (LlamaRMSNorm, "RMSNorm", False),
+        (RMSNorm, "RMSNorm", False),
         (GemmaRMSNorm, "GemmaRMSNorm", True),
         (WeightlessRMSNorm, "RMSNorm", False),
     ],
@@ -542,6 +531,26 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c
     assert built.variance_epsilon == fuser.eps
 
 
+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 fuser.has_weight and not fuser.zero_centered
+        assert fuser.eps is None  # default eps, resolved at fuse time
+        model_config = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32)
+        built = fuser.fuse(module, "norm", model_config, None)
+    assert type(built) is VLLMRMSNorm
+    assert built.variance_epsilon == torch.finfo(torch.float32).eps
+
+
+# MoE fuser
+
+
 class TopKRouter(nn.Module):
     """HF v5 top-k router: `linear -> softmax -> topk (-> renorm)`."""
 

From b6b12dd34a800fd294fa60a35082d5382652be42 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 10:30:32 +0000
Subject: [PATCH 16/35] Concise RMSNorm comments

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/fusers/rms_norm.py       | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py
index 1255ffc56a99..de6259eed006 100644
--- a/vllm/model_executor/models/transformers/fusers/rms_norm.py
+++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py
@@ -58,20 +58,16 @@ def _is_one_plus(node: object) -> bool:
 
 @dataclass
 class RMSNormFuser(BaseFuser):
-    """Fuses an HF RMSNorm's `x * rsqrt(mean(x**2) + eps)` normalize into vLLM's
-    fused RMSNorm CustomOp, read structurally from the module's dataflow.
-
-    eps and the weight form (plain, zero-centered or absent) come straight from
-    the graph, so detection does not depend on the class name.
-    """
+    """Fuser for RMSNorm patterns, including Gemma-style zero-centered weights."""
 
     eps: float | None
     """`None` only for a fused `rms_norm` op with default eps; resolved in `fuse`."""
     has_weight: bool
+    """Does the norm have a weight?"""
     zero_centered: bool
     """Gemma-style `(1 + weight)` scaling (weight initialised at zero)."""
     source_cls: str
-    """Class name of the HF norm this was matched from (for logging)."""
+    """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"
@@ -79,6 +75,7 @@ def info(self, name: str) -> str:
 
     @classmethod
     def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
+        """Match a graph to the RMSNorm pattern, returning a fuser if found."""
         x = find_node(graph, lambda n: n.op == "placeholder")
         if x is None:
             return None
@@ -137,10 +134,7 @@ def fuse(
         model_config: "ModelConfig",
         quant_config: "QuantizationConfig",
     ) -> nn.Module:
-        """Build vLLM's RMSNorm from the detected parameters.
-
-        The scale keeps its checkpoint name (`weight`), so no weight remapping is
-        needed. `var_hidden_size` has no Transformers equivalent."""
+        """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()

From 9c408baa48c9081b653b43592137bc5e77fab57d Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 11:20:44 +0000
Subject: [PATCH 17/35] renaming

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/test_transformers_fusers.py      | 20 +++----
 .../models/transformers/fusers/moe.py         | 56 +++++++++----------
 .../model_executor/models/transformers/moe.py | 39 ++++++-------
 3 files changed, 56 insertions(+), 59 deletions(-)

diff --git a/tests/models/test_transformers_fusers.py b/tests/models/test_transformers_fusers.py
index 14c02498ca30..6c8adc54e7bd 100644
--- a/tests/models/test_transformers_fusers.py
+++ b/tests/models/test_transformers_fusers.py
@@ -19,7 +19,7 @@
     RMSNormFuser,
     get_fuser,
 )
-from vllm.model_executor.models.transformers.fusers.moe import MoEFuser
+from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
 
 # GLU and QKV fusers (shared module fixtures and fuser-application helpers)
 
@@ -744,8 +744,8 @@ def forward(self, hidden_states):
 def test_moe_fuser_detects_router(sigmoid):
     with torch.device("meta"):
         block = MoEBlock(lambda: TopKRouter(sigmoid=sigmoid))
-    fuser = MoEFuser.match(block)
-    assert isinstance(fuser, MoEFuser)
+    fuser = MoEBlockFuser.match(block)
+    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
@@ -754,8 +754,8 @@ def test_moe_fuser_detects_router(sigmoid):
 def test_moe_fuser_detects_shared_experts():
     with torch.device("meta"):
         block = MoEBlockShared()
-    fuser = MoEFuser.match(block)
-    assert isinstance(fuser, MoEFuser)
+    fuser = MoEBlockFuser.match(block)
+    assert isinstance(fuser, MoEBlockFuser)
     assert fuser.shared_name == "shared_expert"
     assert fuser.shared_gate_name == "shared_expert_gate"
     # The shared expert's gate/up projections are merged, scoped to its qualname.
@@ -774,8 +774,8 @@ def test_moe_fuser_detects_shared_experts():
 def test_moe_fuser_shared_without_gate():
     with torch.device("meta"):
         block = MoEBlockSharedNoGate()
-    fuser = MoEFuser.match(block)
-    assert isinstance(fuser, MoEFuser)
+    fuser = MoEBlockFuser.match(block)
+    assert isinstance(fuser, MoEBlockFuser)
     assert fuser.shared_name == "shared_expert"
     assert fuser.shared_gate_name is None
 
@@ -783,8 +783,8 @@ def test_moe_fuser_shared_without_gate():
 def test_moe_fuser_detects_non_glu_shared_expert():
     with torch.device("meta"):
         block = MoEBlockSharedNonGLU()
-    fuser = MoEFuser.match(block)
-    assert isinstance(fuser, MoEFuser)
+    fuser = MoEBlockFuser.match(block)
+    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
@@ -806,4 +806,4 @@ def test_moe_fuser_detects_non_glu_shared_expert():
 def test_moe_fuser_declines_unsupported(block_cls):
     with torch.device("meta"):
         block = block_cls()
-    assert MoEFuser.match(block) is None
+    assert MoEBlockFuser.match(block) is None
diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py
index 15a58b69bc37..8c6bf9634da7 100644
--- a/vllm/model_executor/models/transformers/fusers/moe.py
+++ b/vllm/model_executor/models/transformers/fusers/moe.py
@@ -90,7 +90,7 @@ def _reaches(node: fx.Node, key: str) -> set[fx.Node]:
     return seen
 
 
-def _match_shared_expert(
+def _match_shared_experts(
     graph: fx.Graph, experts: str = "experts"
 ) -> tuple[str | None, str | None]:
     """Detects the shared expert and its optional gate by dataflow."""
@@ -129,13 +129,13 @@ def _match_shared_expert(
 class SharedExpertMLP(nn.Module):
     """Wraps an HF shared expert, applying the output gating it is paired with."""
 
-    def __init__(self, shared_expert: nn.Module, gate: nn.Module | None = None):
+    def __init__(self, shared_experts: nn.Module, gate: nn.Module | None = None):
         super().__init__()
-        self.shared_expert = shared_expert
+        self.shared_experts = shared_experts
         self.gate = gate
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        out = self.shared_expert(hidden_states)
+        out = self.shared_experts(hidden_states)
         if self.gate is not None:
             out = torch.sigmoid(self.gate(hidden_states)[0]) * out
         return out
@@ -160,9 +160,8 @@ def _moe_block_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Te
 
 
 @dataclass
-class MoEFuser:
-    """Fuser for an HF MoE block, routing through vLLM's `FusedMoE` with the
-    same shared expert (if any) and gate (if any) as the original block."""
+class MoEBlockFuser:
+    """Fuser for MoE block `experts`, `gate` and `shared_experts` (optional)."""
 
     gate_name: str
     scoring_func: str
@@ -171,63 +170,60 @@ class MoEFuser:
     shared_glu: GLUFuser | None
 
     @classmethod
-    def match(cls, mlp: nn.Module) -> "MoEFuser | None":
-        # The native block forward returns a single tensor.
-        if _returns_tuple(type(mlp)):
+    def match(cls, moe_block: nn.Module) -> "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 (traced per child).
         gate_name = scoring_func = None
-        for name, child in mlp.named_children():
+        for name, child in moe_block.named_children():
             if name != "experts" and (func := _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: whatever the block adds to the experts' output, found by
-        # dataflow so any MLP works, not just a GLU. A GLU inside it is fused by
-        # the base pass when it descends into the block child; we only precompute
-        # its gate/up remap (see `orig_to_new_stacked`).
-        graph = trace(mlp)
+        # Shared expert: whatever the block adds to the experts' output.
+        graph = trace(moe_block)
         if graph is None:
             return None
-        shared_name, shared_gate_name = _match_shared_expert(graph)
+        shared_name, shared_gate_name = _match_shared_experts(graph)
         if shared_gate_name is not None and not _is_scalar_gate(
-            getattr(mlp, shared_gate_name)
+            getattr(moe_block, shared_gate_name)
         ):
             return None
         shared_glu = None
         if (
             shared_name is not None
-            and (sgraph := trace(shared := getattr(mlp, shared_name))) is not None
+            and (sgraph := trace(shared := getattr(moe_block, shared_name))) is not None
         ):
             shared_glu = GLUFuser.match(sgraph, shared)
         # Fail closed: `rewrite_forward` runs only the experts and the detected
         # shared expert, so any other stateful child would be dropped.
         accounted = {"experts", gate_name, shared_name, shared_gate_name}
-        for name, child in mlp.named_children():
+        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, shared_glu)
 
-    def build_gate(self, mlp: nn.Module, prefix: str) -> ReplicatedLinear:
+    def build_gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear:
         """Rebuild the HF router as a `ReplicatedLinear` producing logits.
 
         Shapes come from the original router weight (`[num_experts, hidden]`),
         the plain `F.linear` weight it loads by identity into. Left unquantized
         -- routers run in full precision -- matching the checkpoint's weight.
         """
-        num_experts, hidden_size = getattr(mlp, self.gate_name).weight.shape
+        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(mlp, self.gate_name, gate)
+        setattr(moe_block, self.gate_name, gate)
         return gate
 
     def build_shared_experts(
-        self, mlp: nn.Module, prefix: str
+        self, moe_block: nn.Module, prefix: str
     ) -> SharedExpertMLP | None:
         """Wrap the HF shared expert, folding in its sigmoid output gate.
 
@@ -238,18 +234,18 @@ def build_shared_experts(
         """
         if self.shared_name is None:
             return None
-        shared_expert = getattr(mlp, self.shared_name)
+        shared_experts = getattr(moe_block, self.shared_name)
         gate = None
         if self.shared_gate_name is not None:
-            hf_gate = getattr(mlp, self.shared_gate_name)
+            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(mlp, self.shared_gate_name, gate)
-        return SharedExpertMLP(shared_expert, gate)
+            setattr(moe_block, self.shared_gate_name, gate)
+        return SharedExpertMLP(shared_experts, gate)
 
     def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
         """Merge the shared expert's gate/up projections (scoped to its qualname).
@@ -263,6 +259,6 @@ def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
             maybe_prefix(prefix, self.shared_name)
         )
 
-    def rewrite_forward(self, mlp: nn.Module) -> None:
+    def rewrite_forward(self, moe_block: nn.Module) -> None:
         """Bind the native block forward (`experts` routes + runs shared)."""
-        mlp.forward = types.MethodType(_moe_block_forward, mlp)
+        moe_block.forward = types.MethodType(_moe_block_forward, moe_block)
diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py
index a309aed3bd9c..7b321c2984ce 100644
--- a/vllm/model_executor/models/transformers/moe.py
+++ b/vllm/model_executor/models/transformers/moe.py
@@ -30,7 +30,7 @@
 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 MoEFuser
+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
 
@@ -146,11 +146,11 @@ def update_physical_experts_metadata(
         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."""
@@ -223,10 +223,10 @@ def _recursive_replace(module: nn.Module, prefix: str):
                 is_3d = len(params) > 0 and all(p.ndim == 3 for p in params)
                 if child_name == "experts" 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)
-                    mlp_cls = type(mlp).__name__
+                    moe_block_cls = type(moe_block).__name__
                     experts_cls = type(experts).__name__
                     # Do the experts have biases
                     has_bias = False
@@ -237,8 +237,8 @@ def _recursive_replace(module: nn.Module, prefix: str):
                     # 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
 
@@ -257,11 +257,12 @@ def _recursive_replace(module: nn.Module, prefix: str):
                         has_bias=has_bias,
                         routed_experts_cls=TransformersRoutedExperts,
                     )
-                    if self.num_expert_groups <= 1 and (fuser := MoEFuser.match(mlp)):
-                        # MLP forward is fully replaced.
+                    fuser = MoEBlockFuser.match(moe_block)
+                    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.
-                        gate = fuser.build_gate(mlp, prefix)
-                        shared_experts = fuser.build_shared_experts(mlp, prefix)
+                        gate = fuser.build_gate(moe_block, prefix)
+                        shared_experts = fuser.build_shared_experts(moe_block, prefix)
                         if shared_experts is not None:
                             self.hf_to_vllm_mapper.orig_to_new_stacked.update(
                                 fuser.orig_to_new_stacked(prefix)
@@ -274,17 +275,17 @@ def _recursive_replace(module: nn.Module, prefix: str):
                             gate=gate,
                             shared_experts=shared_experts,
                         )
-                        fuser.rewrite_forward(mlp)
+                        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,
-                            mlp_cls,
+                            moe_block_cls,
                         )
                     else:
-                        # MLP forward is unmodified.
+                        # 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()
@@ -325,10 +326,10 @@ def custom_routing_function(
                             experts_cls,
                         )
                     fused_experts = FusedMoE(**kwargs)
-                    mlp.experts = fused_experts
+                    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)

From 126e5df1e815375466fdf989a2305b81827540b8 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 13:21:38 +0000
Subject: [PATCH 18/35] Consistency pass

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/test_transformers_fusers.py      | 164 +++++++-----
 .../models/transformers/base.py               |   6 +-
 .../models/transformers/fusers/moe.py         | 250 +++++++++---------
 .../model_executor/models/transformers/moe.py |  15 +-
 4 files changed, 231 insertions(+), 204 deletions(-)

diff --git a/tests/models/test_transformers_fusers.py b/tests/models/test_transformers_fusers.py
index 6c8adc54e7bd..a6e32c7fc9fb 100644
--- a/tests/models/test_transformers_fusers.py
+++ b/tests/models/test_transformers_fusers.py
@@ -93,11 +93,7 @@ def forward(self, x):
 
 
 class FakeAttention(nn.Module):
-    """Mimics the structure of a HF v5 attention module under vLLM.
-
-    The forward exercises the patterns the fusion tracer must support: shape unpacking,
-    the dead KV cache branch, and the `attention_interface` call splatting `**kwargs`
-    (resolved to vLLM's interface, which is traced through)."""
+    """HF v5-style attention: shape unpacking, dead KV branch, kwargs interface."""
 
     is_causal = True
 
@@ -458,8 +454,7 @@ def forward(self, x):
 
 
 class LayerNorm(RMSNorm):
-    """An RMSNorm whose class name is *not* `*RMSNorm` (name-independence), and
-    which keeps the input dtype throughout (no float32 upcast)."""
+    """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)
@@ -569,11 +564,7 @@ def forward(self, hidden_states):
 
 
 class CorrectionRouter(nn.Module):
-    """Grouped/noaux router whose score-correction bias is a buffer (as in
-    DeepSeek-V3).
-
-    It is `weight`-only in its parameters, so it is declined via the correction
-    bias read structurally from the graph, not by the parameter check."""
+    """Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined."""
 
     def __init__(self, num_experts=8, hidden=16):
         super().__init__()
@@ -588,10 +579,7 @@ def forward(self, hidden_states):
 
 
 class BiasedRouter(TopKRouter):
-    """A valid top-k router but not `weight`-only (an extra `bias` parameter).
-
-    `build_gate` rebuilds the router as a bias-free `ReplicatedLinear`, so a
-    router carrying any other parameter cannot be reproduced faithfully."""
+    """A valid top-k router but not `weight`-only (extra `bias` param) -> declined."""
 
     def __init__(self):
         super().__init__()
@@ -604,6 +592,16 @@ def forward(self, hidden_states):
         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."""
 
@@ -617,17 +615,29 @@ def forward(self, hidden_states, index, weights):
 
 
 class MoEBlock(nn.Module):
-    """A single-tensor-returning MoE block (Qwen3-style)."""
+    """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])
-        _, weights, index = self.gate(x)
-        return self.experts(x, index, weights).reshape(hidden_states.shape)
+        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):
@@ -638,12 +648,8 @@ def __init__(self):
         self.shared_expert = GLUMLP()
         self.shared_expert_gate = nn.Linear(16, 1, bias=False)
 
-    def forward(self, hidden_states):
-        x = hidden_states.reshape(-1, hidden_states.shape[-1])
-        _, weights, index = self.gate(x)
-        out = self.experts(x, index, weights)
-        out = out + torch.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
-        return out.reshape(hidden_states.shape)
+    def _shared(self, x, logits):
+        return torch.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
 
 
 class MoEBlockSharedNoGate(MoEBlock):
@@ -653,11 +659,8 @@ def __init__(self):
         super().__init__()
         self.shared_expert = GLUMLP()
 
-    def forward(self, hidden_states):
-        x = hidden_states.reshape(-1, hidden_states.shape[-1])
-        _, weights, index = self.gate(x)
-        out = self.experts(x, index, weights) + self.shared_expert(x)
-        return out.reshape(hidden_states.shape)
+    def _shared(self, x, logits):
+        return self.shared_expert(x)
 
 
 class MoEBlockTuple(MoEBlock):
@@ -669,6 +672,29 @@ def forward(self, hidden_states):
         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."""
 
@@ -683,27 +709,18 @@ def forward(self, x):
 
 
 class MoEBlockSharedNonGLU(MoEBlock):
-    """A block whose shared expert is a non-GLU MLP -> detected by dataflow.
-
-    It is added to the experts' output, so it is recognised as the shared
-    expert even though it is not a GLU (no gate/up merge is needed).
-    """
+    """A non-GLU shared expert -> detected by dataflow (no gate/up merge)."""
 
     def __init__(self):
         super().__init__()
         self.shared_expert = PlainMLP()
 
-    def forward(self, hidden_states):
-        x = hidden_states.reshape(-1, hidden_states.shape[-1])
-        _, weights, index = self.gate(x)
-        out = self.experts(x, index, weights) + self.shared_expert(x)
-        return out.reshape(hidden_states.shape)
+    def _shared(self, x, logits):
+        return self.shared_expert(x)
 
 
 class MoEBlockUnaccounted(MoEBlock):
-    """A block with a weight-bearing child outside the experts + shared-expert
-    dataflow (here a pre-router transform) -> must decline, since the rewritten
-    forward would drop it."""
+    """A weight-bearing child outside the fused dataflow (pre-router) -> declined."""
 
     def __init__(self):
         super().__init__()
@@ -726,25 +743,19 @@ def forward(self, x):
         return x * self.scale
 
 
-class MoEBlockUnaccountedBuffer(MoEBlock):
-    """Like `MoEBlockUnaccounted`, but the dropped child holds only a buffer, so
-    the fail-closed check must inspect buffers as well as parameters."""
+class MoEBlockUnaccountedBuffer(MoEBlockUnaccounted):
+    """Like `MoEBlockUnaccounted`, but the extra child holds only a buffer."""
 
     def __init__(self):
         super().__init__()
         self.extra = BufferScale()
 
-    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)
-
 
 @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)
+    fuser = MoEBlockFuser.match(block, "experts")
     assert isinstance(fuser, MoEBlockFuser)
     assert fuser.gate_name == "gate"
     assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax")
@@ -754,27 +765,26 @@ def test_moe_fuser_detects_router(sigmoid):
 def test_moe_fuser_detects_shared_experts():
     with torch.device("meta"):
         block = MoEBlockShared()
-    fuser = MoEBlockFuser.match(block)
+    fuser = MoEBlockFuser.match(block, "experts")
     assert isinstance(fuser, MoEBlockFuser)
     assert fuser.shared_name == "shared_expert"
     assert fuser.shared_gate_name == "shared_expert_gate"
-    # The shared expert's gate/up projections are merged, scoped to its qualname.
-    assert fuser.orig_to_new_stacked("model.layers.0.mlp") == {
-        "model.layers.0.mlp.shared_expert.gate_proj": (
-            "model.layers.0.mlp.shared_expert.gate_up_proj",
-            0,
-        ),
-        "model.layers.0.mlp.shared_expert.up_proj": (
-            "model.layers.0.mlp.shared_expert.gate_up_proj",
-            1,
-        ),
-    }
+
+
+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)
+    fuser = MoEBlockFuser.match(block, "experts")
     assert isinstance(fuser, MoEBlockFuser)
     assert fuser.shared_name == "shared_expert"
     assert fuser.shared_gate_name is None
@@ -783,14 +793,11 @@ def test_moe_fuser_shared_without_gate():
 def test_moe_fuser_detects_non_glu_shared_expert():
     with torch.device("meta"):
         block = MoEBlockSharedNonGLU()
-    fuser = MoEBlockFuser.match(block)
+    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
-    # A non-GLU shared expert needs no gate/up merge; it loads by identity.
-    assert fuser.shared_glu is None
-    assert fuser.orig_to_new_stacked("model.layers.0.mlp") == {}
 
 
 @pytest.mark.parametrize(
@@ -799,6 +806,7 @@ def test_moe_fuser_detects_non_glu_shared_expert():
         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
     ],
@@ -806,4 +814,20 @@ def test_moe_fuser_detects_non_glu_shared_expert():
 def test_moe_fuser_declines_unsupported(block_cls):
     with torch.device("meta"):
         block = block_cls()
-    assert MoEBlockFuser.match(block) is None
+    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
diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index c55cdf871db6..3d1f737058e9 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -42,6 +42,7 @@
     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,
@@ -453,6 +454,7 @@ def recursive_replace(self):
         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)
+            print(f"Registering fusion for {prefix}: {orig_to_new_stacked}")
             self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked)
 
             packed_modules_mapping = fuser.packed_modules_mapping
@@ -504,7 +506,9 @@ def _recursive_replace(module: nn.Module, prefix: str):
                     )
                     logger.info_once(fuser.info(child_name))
                     _recursive_replace(new_module, prefix=qual_name)
-                else:
+                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:
diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py
index 8c6bf9634da7..6a3c7e85d6e5 100644
--- a/vllm/model_executor/models/transformers/fusers/moe.py
+++ b/vllm/model_executor/models/transformers/fusers/moe.py
@@ -15,36 +15,57 @@
 
 from vllm.distributed import tensor_model_parallel_all_gather
 from vllm.model_executor.layers.linear import ReplicatedLinear
-from vllm.model_executor.models.transformers.fusers.glu import GLUFuser
 from vllm.model_executor.models.transformers.fx_utils import (
     find_node,
     is_op,
     peel,
     trace,
 )
-from vllm.model_executor.models.utils import (
-    ShardId,
-    maybe_prefix,
-    sequence_parallel_chunk,
-)
+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 -- its named parameters and buffers together."""
+    """`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:
-    """True if the class's forward returns a tuple (not a single tensor)."""
+    """Does `cls.forward()` return a tuple?"""
     try:
         source = textwrap.dedent(inspect.getsource(inspect.unwrap(cls.forward)))
-        tree = ast.parse(source)
-    except (OSError, SyntaxError, TypeError):
+        forward = ast.parse(source).body[0]
+    except (OSError, SyntaxError, TypeError, IndexError):
         return True
-    return any(
-        isinstance(node, ast.Return) and isinstance(node.value, ast.Tuple)
-        for node in ast.walk(tree)
-    )
+    # 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:
@@ -58,27 +79,8 @@ def _is_scalar_gate(module: nn.Module) -> bool:
     )
 
 
-def _match_router(gate: nn.Module) -> str | None:
-    """Detects a router by dataflow: linear + softmax/sigmoid + top-k."""
-    if [name for name, _ in named_state(gate)] != ["weight"]:
-        return None
-    graph = trace(gate)
-    if graph is None:
-        return None
-    nodes = list(graph.nodes)
-    if not any(is_op(n, "linear") for n in nodes):
-        return None
-    if not any(is_op(n, "topk") for n in nodes):
-        return None
-    softmax = any(is_op(n, "softmax") for n in nodes)
-    sigmoid = any(is_op(n, "sigmoid") for n in nodes)
-    if softmax == sigmoid:  # need exactly one scoring function
-        return None
-    return "softmax" if softmax else "sigmoid"
-
-
 def _reaches(node: fx.Node, key: str) -> set[fx.Node]:
-    """`node` and every fx node connected to it via `key` (`users`/`inputs`)."""
+    """Returns the set of nodes reachable from `node` by following `key` edges."""
     seen: set[fx.Node] = set()
     stack = [node]
     while stack:
@@ -90,42 +92,6 @@ def _reaches(node: fx.Node, key: str) -> set[fx.Node]:
     return seen
 
 
-def _match_shared_experts(
-    graph: fx.Graph, experts: str = "experts"
-) -> 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 node in graph.nodes:
-        if not is_op(node, "add"):
-            continue
-        operands = [a for a in node.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
-
-
 class SharedExpertMLP(nn.Module):
     """Wraps an HF shared expert, applying the output gating it is paired with."""
 
@@ -142,11 +108,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 def _moe_block_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
-    """Native MoE block forward: the internal router routes + runs shared experts.
+    """Standard MoE block forward.
 
-    Under sequence-parallel MoE, tokens are scattered across TP ranks before the
-    experts and all-gathered back after (matching vLLM's native MoE blocks).
-    """
+    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]
@@ -167,51 +131,108 @@ class MoEBlockFuser:
     scoring_func: str
     shared_name: str | None
     shared_gate_name: str | None
-    shared_glu: GLUFuser | 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) -> "MoEBlockFuser | None":
+    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 (traced per child).
+        # Router: the child that scores + top-k selects.
         gate_name = scoring_func = None
         for name, child in moe_block.named_children():
-            if name != "experts" and (func := _match_router(child)) is not None:
+            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: whatever the block adds to the experts' output.
-        graph = trace(moe_block)
-        if graph is None:
-            return None
-        shared_name, shared_gate_name = _match_shared_experts(graph)
-        if shared_gate_name is not None and not _is_scalar_gate(
-            getattr(moe_block, shared_gate_name)
-        ):
-            return None
-        shared_glu = None
-        if (
-            shared_name is not None
-            and (sgraph := trace(shared := getattr(moe_block, shared_name))) is not None
-        ):
-            shared_glu = GLUFuser.match(sgraph, shared)
+        # 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", gate_name, shared_name, shared_gate_name}
+        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, shared_glu)
-
-    def build_gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear:
-        """Rebuild the HF router as a `ReplicatedLinear` producing logits.
+        return cls(gate_name, scoring_func, shared_name, shared_gate_name)
 
-        Shapes come from the original router weight (`[num_experts, hidden]`),
-        the plain `F.linear` weight it loads by identity into. Left unquantized
-        -- routers run in full precision -- matching the checkpoint's weight.
-        """
+    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,
@@ -222,16 +243,11 @@ def build_gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear:
         setattr(moe_block, self.gate_name, gate)
         return gate
 
-    def build_shared_experts(
+    def shared_experts(
         self, moe_block: nn.Module, prefix: str
     ) -> SharedExpertMLP | None:
-        """Wrap the HF shared expert, folding in its sigmoid output gate.
-
-        The MLP is the original module (converted in place by the base pass); the
-        sibling gate is rebuilt as a `ReplicatedLinear` (its `[1, hidden]` weight
-        is used as a plain `F.linear` weight) so the wrapper holds a reference
-        that survives the base pass's replacement, and it loads by identity.
-        """
+        """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)
@@ -247,18 +263,6 @@ def build_shared_experts(
             setattr(moe_block, self.shared_gate_name, gate)
         return SharedExpertMLP(shared_experts, gate)
 
-    def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
-        """Merge the shared expert's gate/up projections (scoped to its qualname).
-
-        Registered here so loading is correct regardless of whether the base pass
-        reaches the shared expert via the block child or the runner alias first.
-        """
-        if self.shared_name is None or self.shared_glu is None:
-            return {}
-        return self.shared_glu.orig_to_new_stacked(
-            maybe_prefix(prefix, self.shared_name)
-        )
-
     def rewrite_forward(self, moe_block: nn.Module) -> None:
-        """Bind the native block forward (`experts` routes + runs shared)."""
+        """Rewrite `moe_block.forward` to route through vLLM's fused MoE."""
         moe_block.forward = types.MethodType(_moe_block_forward, moe_block)
diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py
index 7b321c2984ce..d1f5dba03733 100644
--- a/vllm/model_executor/models/transformers/moe.py
+++ b/vllm/model_executor/models/transformers/moe.py
@@ -154,6 +154,7 @@ def update_physical_experts_metadata(
 
     def recursive_replace(self):
         """Initialize the MoE layers."""
+        experts_name = "experts"
         text_config = self.text_config
 
         # Positional arguments
@@ -221,7 +222,7 @@ def _recursive_replace(module: nn.Module, prefix: str):
                 # 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
                     moe_block = module
                     experts = child_module
@@ -257,23 +258,17 @@ def _recursive_replace(module: nn.Module, prefix: str):
                         has_bias=has_bias,
                         routed_experts_cls=TransformersRoutedExperts,
                     )
-                    fuser = MoEBlockFuser.match(moe_block)
+                    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.
-                        gate = fuser.build_gate(moe_block, prefix)
-                        shared_experts = fuser.build_shared_experts(moe_block, prefix)
-                        if shared_experts is not None:
-                            self.hf_to_vllm_mapper.orig_to_new_stacked.update(
-                                fuser.orig_to_new_stacked(prefix)
-                            )
                         kwargs |= dict(
                             scoring_func=fuser.scoring_func,
                             is_sequence_parallel=(
                                 self.parallel_config.use_sequence_parallel_moe
                             ),
-                            gate=gate,
-                            shared_experts=shared_experts,
+                            gate=fuser.gate(moe_block, prefix),
+                            shared_experts=fuser.shared_experts(moe_block, prefix),
                         )
                         fuser.rewrite_forward(moe_block)
                         routed = "gate + experts"

From 6faa797ec3648da3c506becde4c3590355dc8d65 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 13:47:41 +0000
Subject: [PATCH 19/35] remove print

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 vllm/model_executor/models/transformers/base.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 3d1f737058e9..b7346879a0d9 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -454,7 +454,6 @@ def recursive_replace(self):
         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)
-            print(f"Registering fusion for {prefix}: {orig_to_new_stacked}")
             self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked)
 
             packed_modules_mapping = fuser.packed_modules_mapping

From 409db2dd99d79638f15a250ebc05e3dbb3f51c52 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 14:43:13 +0000
Subject: [PATCH 20/35] Attempt replacement of down_proj and o_proj

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/fusers/glu.py         | 37 ++++++++++++++-----
 .../models/transformers/fusers/qkv.py         | 26 ++++++++++---
 2 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/vllm/model_executor/models/transformers/fusers/glu.py b/vllm/model_executor/models/transformers/fusers/glu.py
index 1536b1df5ff1..951eb35777ae 100644
--- a/vllm/model_executor/models/transformers/fusers/glu.py
+++ b/vllm/model_executor/models/transformers/fusers/glu.py
@@ -19,11 +19,17 @@
 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:
@@ -49,6 +55,7 @@ class GLUFuser(StackedFuser):
     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"
 
@@ -70,16 +77,16 @@ def _is_act_of_gate(cls, node: fx.Node, module: nn.Module) -> bool:
     @classmethod
     def _get_glu_nodes(
         cls, graph: fx.Graph, module: nn.Module
-    ) -> tuple[fx.Node, fx.Node, fx.Node] | None:
+    ) -> tuple[fx.Node, fx.Node, fx.Node, fx.Node] | None:
         """Search graph for the GLU pattern `act(gate(x)) * up(x)`."""
-        for node in graph.nodes:
+        for mul in graph.nodes:
             if (
-                node.op == "call_function"
-                and node.target == operator.mul
-                and len(node.args) == 2
-                and all(isinstance(arg, fx.Node) for arg in node.args)
+                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 = node.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):
@@ -91,7 +98,7 @@ def _get_glu_nodes(
                     and isinstance(x := gate.args[0], fx.Node)
                     and x is up.args[0]
                 ):
-                    return act, gate, up
+                    return act, gate, up, mul
         return None
 
     @staticmethod
@@ -116,7 +123,7 @@ def _get_act_and_mul(cls, act: nn.Module) -> nn.Module:
     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 = glu_nodes
+        act_node, gate_node, up_node, mul_node = glu_nodes
 
         gate = module.get_submodule(gate_node.target)
         up = module.get_submodule(up_node.target)
@@ -124,11 +131,14 @@ def match(cls, graph: fx.Graph, module: nn.Module) -> "GLUFuser | None":
         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
 
@@ -197,3 +207,12 @@ def update_attrs(
         # 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)
diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py
index 5de0bdca405a..2abb389ae59f 100644
--- a/vllm/model_executor/models/transformers/fusers/qkv.py
+++ b/vllm/model_executor/models/transformers/fusers/qkv.py
@@ -19,6 +19,10 @@
     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:
@@ -35,6 +39,7 @@ class QKVFuser(StackedFuser):
     q_name: str
     k_name: str
     v_name: str
+    o_name: str | None
     merged_name: ClassVar[str] = "qkv_proj"
     merged_cls: ClassVar[str] = "QKVParallelLinear"
 
@@ -78,12 +83,12 @@ 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
-        return cls(
-            source_cls=type(module).__name__,
-            q_name=q.target,
-            k_name=k.target,
-            v_name=v.target,
-        )
+        names = dict(q_name=q.target, k_name=k.target, v_name=v.target)
+        predicate = lambda n, c: isinstance(c, nn.Linear) and n not in names.values()
+        others = [n for n, c in module.named_children() if predicate(n, c)]
+        if len(others) == 1:
+            names["o_name"] = others[0]
+        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."""
@@ -199,3 +204,12 @@ def update_attrs(
         # 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)

From df4bb496b275b34f2898777f2c016af30a567d66 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 15:03:33 +0000
Subject: [PATCH 21/35] Make pp_plan optional

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 docs/models/supported_models.md               |  4 +-
 .../models/transformers/base.py               | 48 ++++++++++++-------
 2 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 215fd059f334..76611fea64fd 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -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, vLLM 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 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
 
diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index b7346879a0d9..9ffcd3db0c49 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -360,12 +360,24 @@ def pipeline_parallel(self):
         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]:
@@ -380,10 +392,9 @@ def setter(obj: object, value: object):
 
         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
 
@@ -393,16 +404,16 @@ def setter(obj: object, value: object):
                 "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(
@@ -410,20 +421,20 @@ def setter(obj: object, value: object):
             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.
@@ -436,14 +447,19 @@ def recursive_replace(self):
         - `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`

From d90bcd9d39169723088ea00ad87fbd68d702743d Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 15:04:50 +0000
Subject: [PATCH 22/35] Small fix

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 vllm/model_executor/models/transformers/fusers/qkv.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py
index 2abb389ae59f..f51eb90921b1 100644
--- a/vllm/model_executor/models/transformers/fusers/qkv.py
+++ b/vllm/model_executor/models/transformers/fusers/qkv.py
@@ -86,8 +86,7 @@ def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None":
         names = dict(q_name=q.target, k_name=k.target, v_name=v.target)
         predicate = lambda n, c: isinstance(c, nn.Linear) and n not in names.values()
         others = [n for n, c in module.named_children() if predicate(n, c)]
-        if len(others) == 1:
-            names["o_name"] = others[0]
+        names["o_name"] = others[0] if len(others) == 1 else None
         return cls(source_cls=type(module).__name__, **names)
 
     def update_forward(self, module: nn.Module) -> None:

From a113e2d5387e5f60e1691b106f852ae7e868d3a2 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 15:05:54 +0000
Subject: [PATCH 23/35] Update tests

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/test_transformers_fusers.py | 52 ++++++++++++++++++++++--
 1 file changed, 48 insertions(+), 4 deletions(-)

diff --git a/tests/models/test_transformers_fusers.py b/tests/models/test_transformers_fusers.py
index a6e32c7fc9fb..fa8b3e2a78a3 100644
--- a/tests/models/test_transformers_fusers.py
+++ b/tests/models/test_transformers_fusers.py
@@ -32,16 +32,26 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return F.silu(x[..., :d]) * x[..., d:]
 
 
-class GLUMLP(nn.Module):
-    """`act(gate(x)) * up(x)` — the canonical HF GLU MLP."""
+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.down_proj = nn.Linear(inter, hidden, 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))
 
@@ -160,6 +170,14 @@ def forward(
         return self.o_proj(attn_output.reshape(*input_shape, -1)), None
 
 
+class ExtraProjAttention(FakeAttention):
+    """A second non-qkv linear -> `o_proj` is ambiguous, so `o_name` is None."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.sink_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+
+
 class FakeSelfAttn(nn.Module):
     """Stand-in for the vLLM `Attention` looked up in `attention_instances`."""
 
@@ -234,7 +252,8 @@ def test_detects_and_rewrites_glu(mlp_cls, bias):
         fuser.gate_name,
         fuser.up_name,
         fuser.act_name,
-    ) == ("gate_proj", "up_proj", "act_fn")
+        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.
@@ -255,6 +274,18 @@ def test_detects_and_rewrites_glu(mlp_cls, bias):
     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):
@@ -270,6 +301,7 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
     # 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)
@@ -299,6 +331,18 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
     torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
 
 
+def test_qkv_identifies_output_projection():
+    """The output projection consuming the attention output is identified.
+
+    It is forced to `RowParallelLinear` in `update_attrs` so its sharded input
+    matches the head-sharded qkv; `None` (falling back to `tp_plan`) when the
+    attention has more than one non-qkv linear, so `o_proj` is ambiguous."""
+    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 is None
+
+
 def test_fuser_is_cached_per_class():
     with torch.device("meta"):
         fuser_a = get_fuser(GLUMLP())

From 8eec984d56e80f93f466f435d43f6ee797eea51f Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Fri, 3 Jul 2026 09:01:25 +0000
Subject: [PATCH 24/35] Always register vllm_attention_forward

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/__init__.py           | 34 +++++++++++++++++++
 .../models/transformers/base.py               | 26 --------------
 2 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py
index cb224e5cbc0e..78a12876e662 100644
--- a/vllm/model_executor/models/transformers/__init__.py
+++ b/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.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): ...
diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 9ffcd3db0c49..8429c83c7a4e 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -31,7 +31,6 @@
     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
@@ -81,31 +80,6 @@
 logger = init_logger(__name__)
 
 
-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
-
-
 class Base(
     nn.Module,
     VllmModel,

From c6e811f62a8748ff326f704beed1aa1c5c7ce830 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Fri, 3 Jul 2026 09:03:29 +0000
Subject: [PATCH 25/35] Remove bad typing pattern

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 vllm/model_executor/models/transformers/base.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py
index 8429c83c7a4e..a36ffacad966 100644
--- a/vllm/model_executor/models/transformers/base.py
+++ b/vllm/model_executor/models/transformers/base.py
@@ -74,8 +74,6 @@
     from transformers import PreTrainedModel
 
     from vllm.config import VllmConfig
-else:
-    PreTrainedModel = object
 
 logger = init_logger(__name__)
 
@@ -195,7 +193,7 @@ def _patch_config(self):
         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.
 
@@ -214,7 +212,7 @@ def _get_decoder_cls(self, **kwargs: dict) -> type[PreTrainedModel]:
 
     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,

From e66f5b92d333042d47edeaccfba1896035040cf7 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Fri, 3 Jul 2026 09:20:35 +0000
Subject: [PATCH 26/35] Restructure tests

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .buildkite/test-amd.yaml                      |   8 +-
 .buildkite/test_areas/models_basic.yaml       |   8 +-
 .buildkite/test_areas/models_distributed.yaml |   2 +-
 .github/CODEOWNERS                            |   2 +-
 tests/models/transformers/__init__.py         |   0
 .../fusers/test_linear.py}                    | 442 +-----------------
 tests/models/transformers/fusers/test_moe.py  | 299 ++++++++++++
 .../transformers/fusers/test_rmsnorm.py       | 144 ++++++
 .../test_backend.py}                          |   8 +-
 .../models/transformers/fusers/__init__.py    |  10 +-
 10 files changed, 472 insertions(+), 451 deletions(-)
 create mode 100644 tests/models/transformers/__init__.py
 rename tests/models/{test_transformers_fusers.py => transformers/fusers/test_linear.py} (52%)
 create mode 100644 tests/models/transformers/fusers/test_moe.py
 create mode 100644 tests/models/transformers/fusers/test_rmsnorm.py
 rename tests/models/{test_transformers.py => transformers/test_backend.py} (97%)

diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml
index 3dd75d526452..e04d5a2d14c7 100644
--- a/.buildkite/test-amd.yaml
+++ b/.buildkite/test-amd.yaml
@@ -476,7 +476,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)'
@@ -1447,10 +1447,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  -----------------------------------------------------#
 
@@ -1683,7 +1683,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
diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml
index 52b3330154c0..87a96c3e24d1 100644
--- a/.buildkite/test_areas/models_basic.yaml
+++ b/.buildkite/test_areas/models_basic.yaml
@@ -41,10 +41,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
@@ -60,7 +60,7 @@ steps:
   - vllm/
   - tests/models/test_utils.py
   - tests/models/test_vision.py
-  - tests/models/test_transformers_fusers.py
+  - tests/models/transformers/fusers/
   device: cpu-small
   commands:
-    - pytest -v -s models/test_utils.py models/test_vision.py models/test_transformers_fusers.py
+    - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/
diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml
index b5758c55affa..a3ee7666ed09 100644
--- a/.buildkite/test_areas/models_distributed.yaml
+++ b/.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
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index 8ca6fc22d64b..57166d9d9b78 100644
--- a/.github/CODEOWNERS
+++ b/.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
diff --git a/tests/models/transformers/__init__.py b/tests/models/transformers/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tests/models/test_transformers_fusers.py b/tests/models/transformers/fusers/test_linear.py
similarity index 52%
rename from tests/models/test_transformers_fusers.py
rename to tests/models/transformers/fusers/test_linear.py
index fa8b3e2a78a3..79b810425641 100644
--- a/tests/models/test_transformers_fusers.py
+++ b/tests/models/transformers/fusers/test_linear.py
@@ -1,27 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the Transformers backend's fx fusers (MoE, GLU, QKV, RMSNorm)."""
+"""Unit tests for the Transformers modeling backend's linear fusers."""
 
 import inspect
-import types
-from types import SimpleNamespace
+from types import MethodType, SimpleNamespace
 
 import pytest
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-# Registers the "vllm" attention interface in ALL_ATTENTION_FUNCTIONS
-import vllm.model_executor.models.transformers.base  # noqa: F401
-from vllm.model_executor.models.transformers.fuser import (
-    GLUFuser,
-    QKVFuser,
-    RMSNormFuser,
-    get_fuser,
-)
-from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
-
-# GLU and QKV fusers (shared module fixtures and fuser-application helpers)
+from vllm.model_executor.models.transformers.fuser import get_fuser
+from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser
 
 
 class SiluAndMulStub(nn.Module):
@@ -214,7 +204,7 @@ def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser):
     setattr(module, fuser.act_name, SiluAndMulStub())
     delattr(module, fuser.gate_name)
     delattr(module, fuser.up_name)
-    module.forward = types.MethodType(fuser.fused_forward, module)
+    module.forward = MethodType(fuser.fused_forward, module)
     return module
 
 
@@ -237,7 +227,7 @@ def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser):
     setattr(module, fuser.merged_name, merged)
     for name in (fuser.q_name, fuser.k_name, fuser.v_name):
         delattr(module, name)
-    module.forward = types.MethodType(fuser.fused_forward, module)
+    module.forward = MethodType(fuser.fused_forward, module)
     return module
 
 
@@ -455,423 +445,3 @@ def test_act_and_mul_derived_from_module(default_vllm_config):
     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())
-
-
-# RMSNorm fuser
-
-
-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)
-
-
-@pytest.mark.parametrize(
-    "cls,eps,has_weight,zero_centered",
-    [
-        (RMSNorm, 1e-5, True, False),
-        (GemmaRMSNorm, 1e-6, True, True),
-        (WeightlessRMSNorm, 1e-6, False, False),
-        (LayerNorm, 1e-6, True, False),
-        (torch.nn.RMSNorm, 1e-5, True, False),  # fused `F.rms_norm` op
-    ],
-)
-def test_detects_rms_norm_variants(cls, eps, has_weight, zero_centered):
-    with torch.device("meta"):
-        fuser = get_fuser(cls(16, eps=eps))
-    assert isinstance(fuser, RMSNormFuser)
-    # eps and the weight form are read from the graph, not the class name.
-    assert fuser.eps == eps
-    assert fuser.has_weight == has_weight
-    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,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)
-    types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm}
-    assert type(built) is types_by_name[expected]
-    assert built.variance_epsilon == fuser.eps
-
-
-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 fuser.has_weight and not fuser.zero_centered
-        assert fuser.eps is None  # default eps, resolved at fuse time
-        model_config = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32)
-        built = fuser.fuse(module, "norm", model_config, None)
-    assert type(built) is VLLMRMSNorm
-    assert built.variance_epsilon == torch.finfo(torch.float32).eps
-
-
-# MoE fuser
-
-
-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
diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py
new file mode 100644
index 000000000000..04eadac3f785
--- /dev/null
+++ b/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
diff --git a/tests/models/transformers/fusers/test_rmsnorm.py b/tests/models/transformers/fusers/test_rmsnorm.py
new file mode 100644
index 000000000000..2c3115974ce1
--- /dev/null
+++ b/tests/models/transformers/fusers/test_rmsnorm.py
@@ -0,0 +1,144 @@
+# 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
+
+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)
+
+
+@pytest.mark.parametrize(
+    "cls,eps,has_weight,zero_centered",
+    [
+        (RMSNorm, 1e-5, True, False),
+        (GemmaRMSNorm, 1e-6, True, True),
+        (WeightlessRMSNorm, 1e-6, False, False),
+        (LayerNorm, 1e-6, True, False),
+        (torch.nn.RMSNorm, 1e-5, True, False),  # fused `F.rms_norm` op
+    ],
+)
+def test_detects_rms_norm_variants(cls, eps, has_weight, zero_centered):
+    with torch.device("meta"):
+        fuser = get_fuser(cls(16, eps=eps))
+    assert isinstance(fuser, RMSNormFuser)
+    # eps and the weight form are read from the graph, not the class name.
+    assert fuser.eps == eps
+    assert fuser.has_weight == has_weight
+    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,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 == fuser.eps
+
+
+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 fuser.has_weight and not fuser.zero_centered
+        assert fuser.eps is None  # default eps, resolved at fuse time
+        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
diff --git a/tests/models/test_transformers.py b/tests/models/transformers/test_backend.py
similarity index 97%
rename from tests/models/test_transformers.py
rename to tests/models/transformers/test_backend.py
index 5f3519552ca7..334075e115f9 100644
--- a/tests/models/test_transformers.py
+++ b/tests/models/transformers/test_backend.py
@@ -6,10 +6,10 @@
 
 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
 
 
 def get_model(arch: str) -> str:
diff --git a/vllm/model_executor/models/transformers/fusers/__init__.py b/vllm/model_executor/models/transformers/fusers/__init__.py
index 4fad6befce62..58910b0ecc3c 100644
--- a/vllm/model_executor/models/transformers/fusers/__init__.py
+++ b/vllm/model_executor/models/transformers/fusers/__init__.py
@@ -4,7 +4,15 @@
 
 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", "QKVFuser", "RMSNormFuser"]
+__all__ = [
+    "BaseFuser",
+    "StackedFuser",
+    "GLUFuser",
+    "MoEBlockFuser",
+    "QKVFuser",
+    "RMSNormFuser",
+]

From 0da7252de63bce75ba61ce2a15f228ab813cdac4 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Fri, 3 Jul 2026 10:01:43 +0000
Subject: [PATCH 27/35] Ensure q_norm or k_norm don't interfere with o_proj
 detection

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/fusers/test_linear.py | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py
index 79b810425641..97216e1f1805 100644
--- a/tests/models/transformers/fusers/test_linear.py
+++ b/tests/models/transformers/fusers/test_linear.py
@@ -168,6 +168,41 @@ def __init__(self, **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`."""
 
@@ -331,6 +366,9 @@ def test_qkv_identifies_output_projection():
         assert get_fuser(FakeAttention()).o_name == "o_proj"
         assert get_fuser(ReversedFakeAttention()).o_name == "o_proj"
         assert get_fuser(ExtraProjAttention()).o_name is None
+        # 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():

From a53f28ecab78dd834f805de454ee42435a837eda Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Fri, 3 Jul 2026 10:35:28 +0000
Subject: [PATCH 28/35] Make RMSNorm TP aware

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/fusers/rms_norm.py    | 47 ++++++++++++++++++-
 1 file changed, 45 insertions(+), 2 deletions(-)

diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py
index de6259eed006..5149f568bc12 100644
--- a/vllm/model_executor/models/transformers/fusers/rms_norm.py
+++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py
@@ -8,6 +8,13 @@
 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, is_op, peel
@@ -56,6 +63,42 @@ def _is_one_plus(node: object) -> bool:
     return any(isinstance(a, (int, float)) and a == 1 for a in node.args)
 
 
+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."""
@@ -145,9 +188,9 @@ def fuse(
             dtype = weight.dtype if weight is not None else model_config.dtype
             eps = torch.finfo(dtype).eps
         if self.zero_centered:
-            return GemmaRMSNorm(hidden_size=hidden_size, eps=eps)
+            return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps)
         has_weight = self.has_weight and weight is not None
-        return RMSNorm(
+        return TPAwareRMSNorm(
             hidden_size=hidden_size,
             eps=eps,
             has_weight=has_weight,

From 814481e577b8ecf1d2c6d8e398b6b9e95232a8bb Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Fri, 3 Jul 2026 10:40:27 +0000
Subject: [PATCH 29/35] Update doc

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 docs/models/supported_models.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 76611fea64fd..7874fa412b21 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -172,7 +172,7 @@ class MyConfig(PretrainedConfig):
 - `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, vLLM 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 and parameter-free modules (e.g. rotary embeddings) on every stage
+    - 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
 

From f2ac331ce58a34e25bde41ae267842eb4e37e4bc Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Fri, 3 Jul 2026 11:06:17 +0000
Subject: [PATCH 30/35] Harden RMSNormFuser

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../models/transformers/fusers/rms_norm.py    | 79 ++++++++++++-------
 .../models/transformers/fx_utils.py           | 16 ++++
 2 files changed, 65 insertions(+), 30 deletions(-)

diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py
index 5149f568bc12..829fd9a541c8 100644
--- a/vllm/model_executor/models/transformers/fusers/rms_norm.py
+++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py
@@ -17,7 +17,13 @@
 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, is_op, peel
+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
@@ -63,6 +69,14 @@ def _is_one_plus(node: object) -> bool:
     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."""
 
@@ -103,10 +117,6 @@ class TPAwareGemmaRMSNorm(TPAwareNormMixin, GemmaRMSNorm):
 class RMSNormFuser(BaseFuser):
     """Fuser for RMSNorm patterns, including Gemma-style zero-centered weights."""
 
-    eps: float | None
-    """`None` only for a fused `rms_norm` op with default eps; resolved in `fuse`."""
-    has_weight: bool
-    """Does the norm have a weight?"""
     zero_centered: bool
     """Gemma-style `(1 + weight)` scaling (weight initialised at zero)."""
     source_cls: str
@@ -119,26 +129,22 @@ def info(self, name: str) -> str:
     @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.
-        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
-            weight = args[2] if len(args) > 2 else kwargs.get("weight")
-            eps = args[3] if len(args) > 3 else kwargs.get("eps")
-            return cls(
-                eps=eps if isinstance(eps, (int, float)) else None,
-                has_weight=isinstance(weight, fx.Node),
-                zero_centered=False,
-                source_cls=type(module).__name__,
-            )
+        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.
-        eps = rsqrt = None
+        rsqrt = None
         for node in graph.nodes:
-            if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None:
+            if is_op(node, "rsqrt") and _variance_eps(node, x) is not None:
                 rsqrt = node
                 break
         if rsqrt is None:
@@ -150,22 +156,34 @@ def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
         if normalize is None:
             return None
         # An optional later `weight * normalized` (or `(1 + weight) * normalized`).
-        has_weight = zero_centered = False
+        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)
-                has_weight = True
-                zero_centered = _is_one_plus(weight)
+                tail, zero_centered = node, _is_one_plus(weight)
                 break
-        return cls(
-            eps=eps,
-            has_weight=has_weight,
-            zero_centered=zero_centered,
-            source_cls=type(module).__name__,
-        )
+        # 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
@@ -182,14 +200,15 @@ def fuse(
         hidden_size = (
             weight.size(0) if weight is not None else model_config.get_hidden_size()
         )
-        eps = self.eps
+        graph = trace(module)
+        eps = self._eps_from_graph(graph) if graph is not None else None
         if eps is None:
-            # Could be `None` for native torch `rms_norm`. Match torch behaviour.
+            # 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 = self.has_weight and weight is not None
+        has_weight = weight is not None
         return TPAwareRMSNorm(
             hidden_size=hidden_size,
             eps=eps,
diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py
index 51310aa126b4..0e043941d8ad 100644
--- a/vllm/model_executor/models/transformers/fx_utils.py
+++ b/vllm/model_executor/models/transformers/fx_utils.py
@@ -129,6 +129,22 @@ def recover_forward(cls: type[nn.Module]) -> tuple[ast.FunctionDef, Callable]:
     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=[])

From 1eca4d77aff0606c04e5d57b550509796df9a1ab Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Fri, 3 Jul 2026 11:42:16 +0000
Subject: [PATCH 31/35] Update RMSNorm tests

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/transformers/fusers/__init__.py  |   0
 .../{test_rmsnorm.py => test_rms_norm.py}     | 109 +++++++++++++++---
 2 files changed, 96 insertions(+), 13 deletions(-)
 create mode 100644 tests/models/transformers/fusers/__init__.py
 rename tests/models/transformers/fusers/{test_rmsnorm.py => test_rms_norm.py} (54%)

diff --git a/tests/models/transformers/fusers/__init__.py b/tests/models/transformers/fusers/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tests/models/transformers/fusers/test_rmsnorm.py b/tests/models/transformers/fusers/test_rms_norm.py
similarity index 54%
rename from tests/models/transformers/fusers/test_rmsnorm.py
rename to tests/models/transformers/fusers/test_rms_norm.py
index 2c3115974ce1..6497b98c4ba3 100644
--- a/tests/models/transformers/fusers/test_rmsnorm.py
+++ b/tests/models/transformers/fusers/test_rms_norm.py
@@ -7,6 +7,7 @@
 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
@@ -71,23 +72,50 @@ def forward(self, x):
         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,has_weight,zero_centered",
+    "cls,eps,zero_centered",
     [
-        (RMSNorm, 1e-5, True, False),
-        (GemmaRMSNorm, 1e-6, True, True),
-        (WeightlessRMSNorm, 1e-6, False, False),
-        (LayerNorm, 1e-6, True, False),
-        (torch.nn.RMSNorm, 1e-5, True, False),  # fused `F.rms_norm` op
+        (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, has_weight, zero_centered):
+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)
-    # eps and the weight form are read from the graph, not the class name.
-    assert fuser.eps == eps
-    assert fuser.has_weight == has_weight
     assert fuser.zero_centered == zero_centered
 
 
@@ -98,6 +126,14 @@ def test_non_rms_norms_are_not_matched(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",
     [
@@ -124,7 +160,10 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c
     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 == fuser.eps
+    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):
@@ -136,9 +175,53 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config):
         module = torch.nn.RMSNorm(16)  # forward is a single `F.rms_norm` call
         fuser = get_fuser(module)
         assert isinstance(fuser, RMSNormFuser)
-        assert fuser.has_weight and not fuser.zero_centered
-        assert fuser.eps is None  # default eps, resolved at fuse time
+        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

From def713ceb5dfa454be3670b4f3cbe23ff27259b6 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Fri, 3 Jul 2026 12:36:41 +0000
Subject: [PATCH 32/35] Harden o_proj detection

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 .../model_executor/models/transformers/fusers/qkv.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py
index f51eb90921b1..03f5a0b49e4a 100644
--- a/vllm/model_executor/models/transformers/fusers/qkv.py
+++ b/vllm/model_executor/models/transformers/fusers/qkv.py
@@ -84,9 +84,15 @@ def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None":
             return None
         q, k, v = qkv_nodes
         names = dict(q_name=q.target, k_name=k.target, v_name=v.target)
-        predicate = lambda n, c: isinstance(c, nn.Linear) and n not in names.values()
-        others = [n for n, c in module.named_children() if predicate(n, c)]
-        names["o_name"] = others[0] if len(others) == 1 else None
+        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:

From 6069c68ea46bbede2f917b16cd944f508216b66b Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Fri, 3 Jul 2026 12:57:12 +0000
Subject: [PATCH 33/35] Update tests

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/transformers/fusers/test_linear.py | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py
index 97216e1f1805..3ead599234f5 100644
--- a/tests/models/transformers/fusers/test_linear.py
+++ b/tests/models/transformers/fusers/test_linear.py
@@ -161,7 +161,7 @@ def forward(
 
 
 class ExtraProjAttention(FakeAttention):
-    """A second non-qkv linear -> `o_proj` is ambiguous, so `o_name` is None."""
+    """A second non-qkv linear of a different width -> `o_proj` still found."""
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -357,15 +357,10 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
 
 
 def test_qkv_identifies_output_projection():
-    """The output projection consuming the attention output is identified.
-
-    It is forced to `RowParallelLinear` in `update_attrs` so its sharded input
-    matches the head-sharded qkv; `None` (falling back to `tp_plan`) when the
-    attention has more than one non-qkv linear, so `o_proj` is ambiguous."""
     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 is None
+        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"

From 23b265e6d1e09f97a64bb77875aaff480e9a05d9 Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 6 Jul 2026 11:43:16 +0000
Subject: [PATCH 34/35] Fix tests

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/transformers/test_backend.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py
index 334075e115f9..dc523f2771c1 100644
--- a/tests/models/transformers/test_backend.py
+++ b/tests/models/transformers/test_backend.py
@@ -12,6 +12,12 @@
 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:
     model_info = HF_EXAMPLE_MODELS.get_hf_info(arch)
     model_info.check_transformers_version(on_fail="skip")

From 702e5a2303645e1f6a4b2c3ea0bce5cac35119ba Mon Sep 17 00:00:00 2001
From: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Date: Mon, 6 Jul 2026 11:52:56 +0000
Subject: [PATCH 35/35] can't pickle local function

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tests/models/transformers/test_backend.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py
index dc523f2771c1..a3eea1783f59 100644
--- a/tests/models/transformers/test_backend.py
+++ b/tests/models/transformers/test_backend.py
@@ -24,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],
@@ -34,16 +45,6 @@ def check_implementation(
     num_fused: tuple[int, int] = (1, 1),
     **kwargs,
 ):
-    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
-
     if kwargs_ref is None:
         kwargs_ref = {}
     if kwargs_test is None: