Skip to content

vllm.model_executor.models.transformers.fusers

Concrete fusers for the Transformers modeling backend.

Modules:

  • attention –

    Attention fuser: the module that dispatches to the attention interface.

  • base –

    Base classes for the Transformers backend fusers.

  • glu –

    GLU projection fuser: act(gate(x)) * up(x) -> a fused gate/up linear.

  • merged_column –

    Fuser for parallel linear projections.

  • mla –

    MLA fuser: adapt a Transformers MLA attention module for vLLM's MLA layer.

  • moe –

    MoE fuser: route an HF MoE block through FusedMoE with vLLM's own routing.

  • packed_qkv –

    Packed-QKV fuser: c_attn(x).split((q, kv, kv)) -> a QKVParallelLinear.

  • qkv –

    QKV projection fuser: q(x), k(x), v(x) -> a fused qkv linear + split.

  • rms_norm –

    RMSNorm fuser: detect the norm structurally and swap in vLLM's fused RMSNorm.

Classes:

  • AttentionFuser –

    A module that dispatches through the Transformers attention interface.

  • BaseFuser –

    A detected fusion and how to apply it.

  • GLUFuser –

    Fuser for the GLU pattern act(gate(x)) * up(x).

  • MLAFuser –

    Fuser for the MLA attention pattern.

  • MergedColumnParallelFuser –

    Fuser for merging column-parallel linear projections.

  • MoEBlockFuser –

    Fuser for MoE block experts, gate and shared_experts (optional).

  • PackedQKVFuser –

    Fuser for attention with q, k and v packed into one projection.

  • QKVFuser –

    Fuser for the attention QKV pattern q(x), k(x), v(x).

  • RMSNormFuser –

    Fuser for RMSNorm patterns, including Gemma-style zero-centered weights.

  • RewriteFuser –

    A fuser that rewrites the module's forward and rebinds it.

  • StackedFuser –

    A fuser that merges sibling projections into one stacked linear and

AttentionFuser dataclass

Bases: BaseFuser

A module that dispatches through the Transformers attention interface.

Methods:

  • layer_index –

    The layer module computes attention for, if it declares one.

  • scale –

    The softmax scale module passes to the interface, or None.

  • sinks –

    The per-head sink tensor module passes to the interface, or None.

  • validate –

    Whether module will actually dispatch to vLLM.

Attributes:

  • s_aux_expr (expr | None) –

    Source of the s_aux= the module hands the interface, if it hands one.

  • scale_expr (expr | None) –

    Source of the scaling= the module hands the interface, if it hands one.

  • source_cls (str) –

    Class of the HF module that dispatches (for logging).

Source code in vllm/model_executor/models/transformers/fusers/attention.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
@dataclass
class AttentionFuser(BaseFuser):
    """A module that dispatches through the Transformers attention interface."""

    redefines_forward: ClassVar[bool] = False

    source_cls: str
    """Class of the HF module that dispatches (for logging)."""
    scale_expr: ast.expr | None = None
    """Source of the `scaling=` the module hands the interface, if it hands one."""
    s_aux_expr: ast.expr | None = None
    """Source of the `s_aux=` the module hands the interface, if it hands one."""

    def info(self, name: str) -> str:
        return f"Found: {name} ({self.source_cls}) -> attention interface"

    @classmethod
    def match(
        cls, graph: fx.Graph | None, module: nn.Module
    ) -> "AttentionFuser | None":
        if (call := interface_call(type(module))) is None:
            return None
        scaling = [kw.value for kw in call.keywords if kw.arg == "scaling"]
        scale_expr = scaling[0] if len(scaling) == 1 else None
        s_aux = [kw.value for kw in call.keywords if kw.arg == "s_aux"]
        s_aux_expr = s_aux[0] if len(s_aux) == 1 else None
        return cls(
            source_cls=type(module).__name__,
            scale_expr=scale_expr,
            s_aux_expr=s_aux_expr,
        )

    def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
        """Whether `module` will actually dispatch to vLLM."""
        config = getattr(module, "config", None)
        # Only patched in the text config, this excludes attention based mm encoders
        vllm_attn_impls = {VLLM_ATTN_IMPL, VLLM_MLA_ATTN_IMPL}
        return getattr(config, "_attn_implementation", None) in vllm_attn_impls

    def fuse(
        self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
    ) -> nn.Module:
        if (sinks := self.sinks(module)) is not None:
            size = sinks.numel() // vllm_config.parallel_config.tensor_parallel_size
            device = vllm_config.device_config.device
            data = torch.empty(size, dtype=sinks.dtype, device=device)
            sinks_param = nn.Parameter(data, requires_grad=False)
            set_weight_attrs(sinks_param, {"weight_loader": sharded_weight_loader(0)})
            setattr(module, self.s_aux_expr.attr, sinks_param)
        return module

    def layer_index(self, module: nn.Module) -> int | None:
        """The layer `module` computes attention for, if it declares one."""
        layer_idx = getattr(module, "layer_idx", None)
        return layer_idx if isinstance(layer_idx, int) else None

    def scale(self, module: nn.Module) -> float | None:
        """The softmax scale `module` passes to the interface, or `None`."""
        if self.scale_expr is None:
            return None
        scale = _resolve(self.scale_expr, module)
        if not isinstance(scale, (int, float)) or isinstance(scale, bool):
            expression = ast.unparse(self.scale_expr)
            raise ValueError(
                f"Cannot resolve attention scaling expression {expression!r} in "
                f"{type(module).__name__}."
            )
        return float(scale)

    def sinks(self, module: nn.Module) -> nn.Parameter | None:
        """The per-head sink tensor `module` passes to the interface, or `None`."""
        if self.s_aux_expr is None:
            return None
        s_aux = _resolve(self.s_aux_expr, module)
        if not isinstance(s_aux, (nn.Parameter, type(None))):
            expression = ast.unparse(self.s_aux_expr)
            raise ValueError(
                f"Cannot resolve attention s_aux expression {expression!r} in "
                f"{type(module).__name__}."
            )
        return s_aux

s_aux_expr = None class-attribute instance-attribute

Source of the s_aux= the module hands the interface, if it hands one.

scale_expr = None class-attribute instance-attribute

Source of the scaling= the module hands the interface, if it hands one.

source_cls instance-attribute

Class of the HF module that dispatches (for logging).

layer_index(module)

The layer module computes attention for, if it declares one.

Source code in vllm/model_executor/models/transformers/fusers/attention.py 
131
132
133
134
def layer_index(self, module: nn.Module) -> int | None:
    """The layer `module` computes attention for, if it declares one."""
    layer_idx = getattr(module, "layer_idx", None)
    return layer_idx if isinstance(layer_idx, int) else None

scale(module)

The softmax scale module passes to the interface, or None.

Source code in vllm/model_executor/models/transformers/fusers/attention.py 
136
137
138
139
140
141
142
143
144
145
146
147
def scale(self, module: nn.Module) -> float | None:
    """The softmax scale `module` passes to the interface, or `None`."""
    if self.scale_expr is None:
        return None
    scale = _resolve(self.scale_expr, module)
    if not isinstance(scale, (int, float)) or isinstance(scale, bool):
        expression = ast.unparse(self.scale_expr)
        raise ValueError(
            f"Cannot resolve attention scaling expression {expression!r} in "
            f"{type(module).__name__}."
        )
    return float(scale)

sinks(module)

The per-head sink tensor module passes to the interface, or None.

Source code in vllm/model_executor/models/transformers/fusers/attention.py 
149
150
151
152
153
154
155
156
157
158
159
160
def sinks(self, module: nn.Module) -> nn.Parameter | None:
    """The per-head sink tensor `module` passes to the interface, or `None`."""
    if self.s_aux_expr is None:
        return None
    s_aux = _resolve(self.s_aux_expr, module)
    if not isinstance(s_aux, (nn.Parameter, type(None))):
        expression = ast.unparse(self.s_aux_expr)
        raise ValueError(
            f"Cannot resolve attention s_aux expression {expression!r} in "
            f"{type(module).__name__}."
        )
    return s_aux

validate(module, vllm_config)

Whether module will actually dispatch to vLLM.

Source code in vllm/model_executor/models/transformers/fusers/attention.py 
112
113
114
115
116
117
def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
    """Whether `module` will actually dispatch to vLLM."""
    config = getattr(module, "config", None)
    # Only patched in the text config, this excludes attention based mm encoders
    vllm_attn_impls = {VLLM_ATTN_IMPL, VLLM_MLA_ATTN_IMPL}
    return getattr(config, "_attn_implementation", None) in vllm_attn_impls

BaseFuser dataclass

Bases: ABC

A detected fusion and how to apply it.

match analyses the module class once (cached, see get_fusers); fuse then applies the fusion to an instance in recursive_replace, returning the module to install in its place.

Methods:

  • fuse –

    Apply the fusion to an already-validated module, returning the

  • info –

    A human-readable description of the fusion at name, for logging.

  • match –

    Match the pattern in graph, returning a fuser if found.

  • orig_to_new_stacked –

    WeightsMapper.orig_to_new_stacked entries this fuser contributes

  • validate –

    Whether this fuser can be applied to this module instance.

Attributes:

Source code in vllm/model_executor/models/transformers/fusers/base.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
@dataclass
class BaseFuser(ABC):
    """A detected fusion and how to apply it.

    `match` analyses the module *class* once (cached, see `get_fusers`); `fuse`
    then applies the fusion to an instance in `recursive_replace`, returning the
    module to install in its place.
    """

    redefines_forward: ClassVar[bool] = True
    """Whether `fuse` gives the module a different forward,
    by rewriting its source or by returning a different module in its place."""

    @abstractmethod
    def info(self, name: str) -> str:
        """A human-readable description of the fusion at `name`, for logging."""

    @classmethod
    @abstractmethod
    def match(cls, graph: fx.Graph, module: nn.Module) -> "BaseFuser | None":
        """Match the pattern in `graph`, returning a fuser if found."""

    @abstractmethod
    def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
        """Whether this fuser can be applied to this `module` instance."""

    @abstractmethod
    def fuse(
        self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
    ) -> 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 {}

packed_modules_mapping property

packed_modules_mapping entries this fuser contributes (none unless it stacks weights).

redefines_forward = True class-attribute

Whether fuse gives the module a different forward, by rewriting its source or by returning a different module in its place.

fuse(module, prefix, vllm_config) abstractmethod

Apply the fusion to an already-validated module, returning the module to install in its place (mutated in place, or freshly built).

Source code in vllm/model_executor/models/transformers/fusers/base.py 
53
54
55
56
57
58
@abstractmethod
def fuse(
    self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
) -> nn.Module:
    """Apply the fusion to an already-validated `module`, returning the
    module to install in its place (mutated in place, or freshly built)."""

info(name) abstractmethod

A human-readable description of the fusion at name, for logging.

Source code in vllm/model_executor/models/transformers/fusers/base.py 
40
41
42
@abstractmethod
def info(self, name: str) -> str:
    """A human-readable description of the fusion at `name`, for logging."""

match(graph, module) abstractmethod classmethod

Match the pattern in graph, returning a fuser if found.

Source code in vllm/model_executor/models/transformers/fusers/base.py 
44
45
46
47
@classmethod
@abstractmethod
def match(cls, graph: fx.Graph, module: nn.Module) -> "BaseFuser | None":
    """Match the pattern in `graph`, returning a fuser if found."""

orig_to_new_stacked(prefix)

WeightsMapper.orig_to_new_stacked entries this fuser contributes (none unless it stacks weights).

Source code in vllm/model_executor/models/transformers/fusers/base.py 
60
61
62
63
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 {}

validate(module, vllm_config) abstractmethod

Whether this fuser can be applied to this module instance.

Source code in vllm/model_executor/models/transformers/fusers/base.py 
49
50
51
@abstractmethod
def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
    """Whether this fuser can be applied to this `module` instance."""

GLUFuser dataclass

Bases: MergedColumnParallelFuser

Fuser for the GLU pattern act(gate(x)) * up(x).

Methods:

  • update_forward –

    Replace act(gate(x)) * up(x) with act(gate_up(x)) in source.

Source code in vllm/model_executor/models/transformers/fusers/glu.py 
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
@dataclass
class GLUFuser(MergedColumnParallelFuser):
    """Fuser for the GLU pattern `act(gate(x)) * up(x)`."""

    act_name: str
    down_name: str | None
    merged_name: ClassVar[str] = "gate_up_proj"

    @property
    def gate_name(self) -> str:
        return self.linear_names[0]

    @property
    def up_name(self) -> str:
        return self.linear_names[1]

    @classmethod
    def _is_act_of_gate(cls, node: fx.Node, module: nn.Module) -> bool:
        """Is node `act(gate(x))` where `gate` is linear and `act` is not linear."""
        return (
            node.op == "call_module"
            and not is_linear(node, module)
            and len(node.args) == 1
            and isinstance(node.args[0], fx.Node)
            and is_linear(node.args[0], module)
        )

    @classmethod
    def _get_glu_nodes(
        cls, graph: fx.Graph, module: nn.Module
    ) -> tuple[fx.Node, fx.Node, fx.Node, fx.Node] | None:
        """Search graph for the GLU pattern `act(gate(x)) * up(x)`."""
        for mul in graph.nodes:
            if (
                mul.op == "call_function"
                and mul.target == operator.mul
                and len(mul.args) == 2
                and all(isinstance(arg, fx.Node) for arg in mul.args)
            ):
                a, b = mul.args
                if cls._is_act_of_gate(a, module) and is_linear(b, module):
                    act, gate, up = a, a.args[0], b
                elif cls._is_act_of_gate(b, module) and is_linear(a, module):
                    act, gate, up = b, b.args[0], a
                else:
                    continue
                if (
                    all(len(args) == 1 for args in (gate.args, up.args))
                    and isinstance(x := gate.args[0], fx.Node)
                    and x is up.args[0]
                ):
                    return act, gate, up, mul
        return None

    @staticmethod
    def _get_act_and_mul_name(act: nn.Module) -> str | None:
        """Get the name of `act` if it has an `...AndMul` equivalent."""
        for name in CLS2ACT.get(type(act), []):
            if name in ACT_AND_MUL_NAMES:
                return name
        # nn.GELU is not in ACT2CLS, but could be in model code
        if type(act) is nn.GELU:
            return "gelu_pytorch_tanh" if act.approximate == "tanh" else "gelu"
        return None

    @classmethod
    def _get_act_and_mul(cls, act: nn.Module) -> nn.Module:
        """Get the `...AndMul` equivalent of a Transformers activation module."""
        if name := cls._get_act_and_mul_name(act):
            return get_act_and_mul_fn(name)
        raise ValueError(f"No AndMul equivalent for {type(act)}")

    @classmethod
    def match(cls, graph: fx.Graph, module: nn.Module) -> "GLUFuser | None":
        if (glu_nodes := cls._get_glu_nodes(graph, module)) is None:
            return None
        act_node, gate_node, up_node, mul_node = glu_nodes

        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,
            linear_names=(gate_node.target, up_node.target),
            down_name=down_node.target if down_node is not None else 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)
        calls = self._unguarded_calls(funcdef, (self.gate_name, self.up_name))
        gate_call, up_call = calls
        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, vllm_config: "VllmConfig") -> bool:
        if not super().validate(module, vllm_config):
            return False
        # An AndMul kernel splits its input in half, so gate and up must match.
        gate = module.get_submodule(self.gate_name)
        up = module.get_submodule(self.up_name)
        if gate.out_features != up.out_features:
            logger.debug("%s and %s differ in size; skipping fusion", gate, up)
            return False
        act = module.get_submodule(self.act_name)
        if self._get_act_and_mul_name(act) is None:
            logger.debug("%s has no AndMul equivalent; skipping fusion", act)
            return False
        return True

    def update_attrs(
        self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
    ) -> None:
        super().update_attrs(module, prefix, vllm_config)
        act_fn = self._get_act_and_mul(module.get_submodule(self.act_name))
        setattr(module, self.act_name, act_fn)
        # 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", vllm_config.quant_config, prefix=down_prefix
            )
            setattr(module, self.down_name, new_down)
            log_replacement(down_prefix, down, new_down)

_get_act_and_mul(act) classmethod

Get the ...AndMul equivalent of a Transformers activation module.

Source code in vllm/model_executor/models/transformers/fusers/glu.py 
114
115
116
117
118
119
@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)}")

_get_act_and_mul_name(act) staticmethod

Get the name of act if it has an ...AndMul equivalent.

Source code in vllm/model_executor/models/transformers/fusers/glu.py 
103
104
105
106
107
108
109
110
111
112
@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

_get_glu_nodes(graph, module) classmethod

Search graph for the GLU pattern act(gate(x)) * up(x).

Source code in vllm/model_executor/models/transformers/fusers/glu.py 
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
@classmethod
def _get_glu_nodes(
    cls, graph: fx.Graph, module: nn.Module
) -> tuple[fx.Node, fx.Node, fx.Node, fx.Node] | None:
    """Search graph for the GLU pattern `act(gate(x)) * up(x)`."""
    for mul in graph.nodes:
        if (
            mul.op == "call_function"
            and mul.target == operator.mul
            and len(mul.args) == 2
            and all(isinstance(arg, fx.Node) for arg in mul.args)
        ):
            a, b = mul.args
            if cls._is_act_of_gate(a, module) and is_linear(b, module):
                act, gate, up = a, a.args[0], b
            elif cls._is_act_of_gate(b, module) and is_linear(a, module):
                act, gate, up = b, b.args[0], a
            else:
                continue
            if (
                all(len(args) == 1 for args in (gate.args, up.args))
                and isinstance(x := gate.args[0], fx.Node)
                and x is up.args[0]
            ):
                return act, gate, up, mul
    return None

_is_act_of_gate(node, module) classmethod

Is node act(gate(x)) where gate is linear and act is not linear.

Source code in vllm/model_executor/models/transformers/fusers/glu.py 
65
66
67
68
69
70
71
72
73
74
@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)
    )

update_forward(module)

Replace act(gate(x)) * up(x) with act(gate_up(x)) in source.

Source code in vllm/model_executor/models/transformers/fusers/glu.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
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)
    calls = self._unguarded_calls(funcdef, (self.gate_name, self.up_name))
    gate_call, up_call = calls
    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)

MLAFuser dataclass

Bases: StackedFuser

Fuser for the MLA attention pattern.

Methods:

  • update_forward –

    Merge q_a_proj and kv_a_proj into one fused down proj then split.

Attributes:

  • shards (list[tuple[str, ShardId]]) –

    q_a_proj and kv_a_proj_with_mqa stack into one down-projection.

Source code in vllm/model_executor/models/transformers/fusers/mla.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
@dataclass
class MLAFuser(StackedFuser):
    """Fuser for the MLA attention pattern."""

    q_proj_name: str | None
    q_a_proj_name: str | None
    q_a_layernorm_name: str | None
    q_b_proj_name: str | None
    kv_a_proj_name: str
    kv_a_layernorm_name: str
    kv_b_proj_name: str
    o_proj_name: str | None
    merged_name: ClassVar[str] = "fused_qkv_a_proj"
    merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear"

    @property
    def has_q_lora(self) -> bool:
        return self.q_a_proj_name is not None

    def info(self, name: str) -> str:
        info_str = (
            f"Fused: {name} ({self.source_cls}) -> MLAAttention (attention interface)"
        )
        if self.has_q_lora:
            info_str += "; " + super().info(name).removeprefix("Fused: ")
        return info_str

    @property
    def shards(self) -> list[tuple[str, ShardId]]:
        """`q_a_proj` and `kv_a_proj_with_mqa` stack into one down-projection."""
        if self.has_q_lora:
            return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)]
        return []

    @property
    def packed_modules_mapping(self) -> dict[str, list[str]]:
        if self.has_q_lora:
            return super().packed_modules_mapping
        return {}

    @classmethod
    def match(cls, graph: fx.Graph, module: nn.Module) -> "MLAFuser | None":
        # Find all `rms_norm(linear(placeholder))` chains.
        chains = []
        for node in graph.nodes:
            if node.op != "call_module" or is_linear(node, module) or not node.args:
                continue
            source = upstream_linear(node.args[0], module)
            if source is None or not _consumes_placeholder(source):
                continue
            if not _is_rms_norm(module.get_submodule(node.target)):
                continue
            chains.append((source, node))

        # Tell the chains apart by width.
        def is_kv_a(chain) -> bool:
            source, rms_norm = chain
            source_mod = module.get_submodule(source.target)
            rms_norm_mod = module.get_submodule(rms_norm.target)
            return source_mod.out_features != _norm_size(rms_norm_mod)

        kv_a_chains = [chain for chain in chains if is_kv_a(chain)]
        q_a_chains = [chain for chain in chains if not is_kv_a(chain)]
        # Exactly one KV chain is MLA's signature. The query chain is optional.
        if len(kv_a_chains) != 1 or len(q_a_chains) > 1:
            return None
        kv_a_proj, kv_a_layernorm = kv_a_chains[0]

        # Linear children claimed for a role so far; the rest resolve by elimination.
        claimed_linears = {kv_a_proj.target}
        q_proj_name = q_a_proj = q_a_layernorm = q_b_proj = None
        if q_a_chains:
            # Find `q_b_proj(q_a_layernorm(q_a_proj(placeholder)))`.
            q_a_proj, q_a_layernorm = q_a_chains[0]
            q_b_proj = downstream_linear(q_a_layernorm, module)
            if q_b_proj is None:
                return None
            claimed_linears |= {q_a_proj.target, q_b_proj.target}
        else:
            # Find `q_proj(placeholder)`.
            placeholder_linears = {
                node.target
                for node in graph.nodes
                if is_linear(node, module) and _consumes_placeholder(node)
            }
            if len(q_proj_candidates := placeholder_linears - claimed_linears) != 1:
                return None
            q_proj_name = next(iter(q_proj_candidates))
            claimed_linears.add(q_proj_name)

        # Find `kv_b_proj(kv_a_layernorm(...))`.
        kv_b_proj = downstream_linear(kv_a_layernorm, module)
        if kv_b_proj is None or kv_b_proj.target in claimed_linears:
            return None
        claimed_linears.add(kv_b_proj.target)

        # Find `o_proj` if it is returned by the forward graph.
        o_proj_name = returned_linear(graph, module)
        if o_proj_name in claimed_linears:
            o_proj_name = None

        return cls(
            source_cls=type(module).__name__,
            q_proj_name=q_proj_name,
            q_a_proj_name=q_a_proj.target if q_a_proj else None,
            q_a_layernorm_name=q_a_layernorm.target if q_a_layernorm else None,
            q_b_proj_name=q_b_proj.target if q_b_proj else None,
            kv_a_proj_name=kv_a_proj.target,
            kv_a_layernorm_name=kv_a_layernorm.target,
            kv_b_proj_name=kv_b_proj.target,
            o_proj_name=o_proj_name,
        )

    def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
        return vllm_config.model_config.use_mla

    def update_forward(self, module: nn.Module) -> None:
        """Merge `q_a_proj` and `kv_a_proj` into one fused down proj then split.
        Bypass the KV expansion method so the compressed latent reaches the `vllm_mla`
        attention interface unexpanded."""
        funcdef, fn = recover_forward(type(module))
        if self.has_q_lora:
            # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`.
            # The fused call is inserted at the top-level statement preceding both.
            names = [self.q_a_proj_name, self.kv_a_proj_name]
            calls = self._unguarded_calls(funcdef, names)
            if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]):
                raise ValueError("down-projections read different inputs")
            indices = [_top_level_index(funcdef, call) for call in calls]
            self._check_input_stable(funcdef, module, calls, funcdef.body, indices)
            self._splice_merged_split(funcdef, calls, funcdef.body, min(indices))

        # Transformers expands the latent into full key/value in a dedicated method.
        # `MLAAttention` consumes the latent directly (absorbing `kv_b_proj`),
        # so replace the expansion call with its own arguments so `kv_c_normed, k_pe`
        # flow to the interface in place of `key, value`.
        expand_call = _single_expand_call(funcdef, module, self.kv_b_proj_name)
        replace_expr(
            funcdef, expand_call, ast.Tuple(elts=list(expand_call.args), ctx=ast.Load())
        )
        self.fused_forward = compile_forward(funcdef, fn)

    def update_attrs(self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"):
        quant_config = vllm_config.quant_config

        def replace_linear_by_name(name: str, style: str):
            linear = module.get_submodule(name)
            replacement = replace_linear_class(
                linear, style, quant_config, prefix=maybe_prefix(prefix, name)
            )
            setattr(module, name, replacement)
            log_replacement(maybe_prefix(prefix, name), linear, replacement)

        if self.has_q_lora:
            q_a = module.get_submodule(self.q_a_proj_name)
            kv_a = module.get_submodule(self.kv_a_proj_name)
            merged = MergedColumnParallelLinear(
                input_size=q_a.in_features,
                output_sizes=[q_a.out_features, kv_a.out_features],
                bias=q_a.bias is not None,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, self.merged_name),
                return_bias=False,
                disable_tp=True,
            )
            logger.debug(
                "%s: %s, %s: %s -> %s: %s",
                self.q_a_proj_name,
                q_a,
                self.kv_a_proj_name,
                kv_a,
                self.merged_name,
                merged,
            )
            setattr(module, self.merged_name, merged)
            # The rewritten forward calls the merged projection instead.
            delattr(module, self.q_a_proj_name)
            delattr(module, self.kv_a_proj_name)
            replace_linear_by_name(self.q_b_proj_name, "colwise")
        else:
            replace_linear_by_name(self.kv_a_proj_name, "replicate")
            replace_linear_by_name(self.q_proj_name, "colwise")

        replace_linear_by_name(self.kv_b_proj_name, "colwise")
        # MLAAttention calls kv_b_proj and expects vLLM's default return_bias=True
        module.get_submodule(self.kv_b_proj_name).return_bias = True
        if self.o_proj_name is not None:
            replace_linear_by_name(self.o_proj_name, "rowwise")

shards property

q_a_proj and kv_a_proj_with_mqa stack into one down-projection.

update_forward(module)

Merge q_a_proj and kv_a_proj into one fused down proj then split. Bypass the KV expansion method so the compressed latent reaches the vllm_mla attention interface unexpanded.

Source code in vllm/model_executor/models/transformers/fusers/mla.py 
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
def update_forward(self, module: nn.Module) -> None:
    """Merge `q_a_proj` and `kv_a_proj` into one fused down proj then split.
    Bypass the KV expansion method so the compressed latent reaches the `vllm_mla`
    attention interface unexpanded."""
    funcdef, fn = recover_forward(type(module))
    if self.has_q_lora:
        # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`.
        # The fused call is inserted at the top-level statement preceding both.
        names = [self.q_a_proj_name, self.kv_a_proj_name]
        calls = self._unguarded_calls(funcdef, names)
        if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]):
            raise ValueError("down-projections read different inputs")
        indices = [_top_level_index(funcdef, call) for call in calls]
        self._check_input_stable(funcdef, module, calls, funcdef.body, indices)
        self._splice_merged_split(funcdef, calls, funcdef.body, min(indices))

    # Transformers expands the latent into full key/value in a dedicated method.
    # `MLAAttention` consumes the latent directly (absorbing `kv_b_proj`),
    # so replace the expansion call with its own arguments so `kv_c_normed, k_pe`
    # flow to the interface in place of `key, value`.
    expand_call = _single_expand_call(funcdef, module, self.kv_b_proj_name)
    replace_expr(
        funcdef, expand_call, ast.Tuple(elts=list(expand_call.args), ctx=ast.Load())
    )
    self.fused_forward = compile_forward(funcdef, fn)

MergedColumnParallelFuser dataclass

Bases: StackedFuser

Fuser for merging column-parallel linear projections.

Methods:

  • match –

    Fuse the module's sibling linears when there is only one such group.

  • update_attrs –

    Replace the module's parallel linears with one merged projection.

  • update_forward –

    Replace the parallel calls with one merged call and split.

  • validate –

    Check that the parallel linears are compatible for merging.

Attributes:

  • merged_name (str) –

    Programmatic name for the merged projection, based on the original names.

Source code in vllm/model_executor/models/transformers/fusers/merged_column.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
@dataclass
class MergedColumnParallelFuser(StackedFuser):
    """Fuser for merging column-parallel linear projections."""

    linear_names: tuple[str, ...]
    merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear"

    @property
    def merged_name(self) -> str:
        """Programmatic name for the merged projection, based on the original names."""
        # len is used to disambiguate names like `a_b` + `c` vs `a` + `b_c`
        parts = "_".join(f"{len(name)}_{name}" for name in self.linear_names)
        return f"merged_proj_{parts}"

    @property
    def shards(self) -> list[tuple[str, ShardId]]:
        return [(name, index) for index, name in enumerate(self.linear_names)]

    @staticmethod
    def _names(group: list[fx.Node]) -> tuple[str, ...] | None:
        names = tuple(str(node.target) for node in group)
        if len(set(names)) != len(names) or any("." in name for name in names):
            return None
        return names

    @classmethod
    def match(
        cls, graph: fx.Graph, module: nn.Module
    ) -> "MergedColumnParallelFuser | None":
        """Fuse the module's sibling linears when there is only one such group."""
        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)
                # Like QKVFuser/PackedQKVFuser: this fuser has no head or
                # TP-replication awareness, so it must not absorb a QKV-shaped
                # pattern (e.g. behind a norm) a more specific fuser declined.
                and node.args[0].op == "placeholder"
            ):
                by_input.setdefault(node.args[0], []).append(node)
        # A group whose members are not distinct direct children is dropped:
        # the source rewrite addresses each projection as `self.<name>` once.
        groups = [
            (group, names)
            for group in by_input.values()
            if len(group) >= 2 and (names := cls._names(group)) is not None
        ]
        if len(groups) > 1:
            logger.debug(
                "%s has %d fusable sibling-linear groups; skipping fusion "
                "since which one to merge is ambiguous",
                type(module).__name__,
                len(groups),
            )
        if len(groups) != 1:
            return None
        _, names = groups[0]
        fuser = cls(source_cls=type(module).__name__, linear_names=names)
        return None if hasattr(module, fuser.merged_name) else fuser

    def update_forward(self, module: nn.Module) -> None:
        """Replace the parallel calls with one merged call and split."""
        funcdef, fn = recover_forward(type(module))
        calls = self._unguarded_calls(funcdef, self.linear_names)
        if len(set(ast.dump(call.args[0]) for call in calls)) != 1:
            raise ValueError("parallel linears read different inputs")
        chains = [block_chain(funcdef.body, call) for call in calls]
        if any(not chain for chain in chains):
            raise ValueError("parallel linear calls not found in the function body")
        blocks = [chain[-1] for chain in chains]
        if len(set(id(block) for block, _ in blocks)) != 1:
            raise ValueError("parallel linear calls are in different blocks")

        block = blocks[0][0]
        indices = [index for _, index in blocks]
        index = min(indices)
        self._check_input_stable(funcdef, module, calls, block, indices)

        # l1(x), l2(x), ... -> merged(x).split(merged.output_sizes / merged.tp_size, -1)
        self._splice_merged_split(funcdef, calls, block, index)
        self.fused_forward = compile_forward(funcdef, fn)

    def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
        """Check that the parallel linears are compatible for merging."""
        linear_layers = [module.get_submodule(name) for name in self.linear_names]
        tp_size = vllm_config.parallel_config.tensor_parallel_size
        return (
            len(linear_layers) >= 2
            and len(set(linear.in_features for linear in linear_layers)) == 1
            and len(set(linear.bias is None for linear in linear_layers)) == 1
            and all(linear.out_features % tp_size == 0 for linear in linear_layers)
        )

    def update_attrs(
        self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
    ) -> None:
        """Replace the module's parallel linears with one merged projection."""
        linear_modules = [module.get_submodule(name) for name in self.linear_names]
        merged = MergedColumnParallelLinear(
            input_size=linear_modules[0].in_features,
            output_sizes=[linear.out_features for linear in linear_modules],
            bias=linear_modules[0].bias is not None,
            quant_config=vllm_config.quant_config,
            prefix=maybe_prefix(prefix, self.merged_name),
            return_bias=False,
        )
        setattr(module, self.merged_name, merged)
        for name in self.linear_names:
            delattr(module, name)
        logger.debug(
            "%s -> %s: %s",
            ", ".join(f"{n}: {m}" for n, m in zip(self.linear_names, linear_modules)),
            self.merged_name,
            merged,
        )

merged_name property

Programmatic name for the merged projection, based on the original names.

match(graph, module) classmethod

Fuse the module's sibling linears when there is only one such group.

Source code in vllm/model_executor/models/transformers/fusers/merged_column.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
@classmethod
def match(
    cls, graph: fx.Graph, module: nn.Module
) -> "MergedColumnParallelFuser | None":
    """Fuse the module's sibling linears when there is only one such group."""
    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)
            # Like QKVFuser/PackedQKVFuser: this fuser has no head or
            # TP-replication awareness, so it must not absorb a QKV-shaped
            # pattern (e.g. behind a norm) a more specific fuser declined.
            and node.args[0].op == "placeholder"
        ):
            by_input.setdefault(node.args[0], []).append(node)
    # A group whose members are not distinct direct children is dropped:
    # the source rewrite addresses each projection as `self.<name>` once.
    groups = [
        (group, names)
        for group in by_input.values()
        if len(group) >= 2 and (names := cls._names(group)) is not None
    ]
    if len(groups) > 1:
        logger.debug(
            "%s has %d fusable sibling-linear groups; skipping fusion "
            "since which one to merge is ambiguous",
            type(module).__name__,
            len(groups),
        )
    if len(groups) != 1:
        return None
    _, names = groups[0]
    fuser = cls(source_cls=type(module).__name__, linear_names=names)
    return None if hasattr(module, fuser.merged_name) else fuser

update_attrs(module, prefix, vllm_config)

Replace the module's parallel linears with one merged projection.

Source code in vllm/model_executor/models/transformers/fusers/merged_column.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def update_attrs(
    self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
) -> None:
    """Replace the module's parallel linears with one merged projection."""
    linear_modules = [module.get_submodule(name) for name in self.linear_names]
    merged = MergedColumnParallelLinear(
        input_size=linear_modules[0].in_features,
        output_sizes=[linear.out_features for linear in linear_modules],
        bias=linear_modules[0].bias is not None,
        quant_config=vllm_config.quant_config,
        prefix=maybe_prefix(prefix, self.merged_name),
        return_bias=False,
    )
    setattr(module, self.merged_name, merged)
    for name in self.linear_names:
        delattr(module, name)
    logger.debug(
        "%s -> %s: %s",
        ", ".join(f"{n}: {m}" for n, m in zip(self.linear_names, linear_modules)),
        self.merged_name,
        merged,
    )

update_forward(module)

Replace the parallel calls with one merged call and split.

Source code in vllm/model_executor/models/transformers/fusers/merged_column.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def update_forward(self, module: nn.Module) -> None:
    """Replace the parallel calls with one merged call and split."""
    funcdef, fn = recover_forward(type(module))
    calls = self._unguarded_calls(funcdef, self.linear_names)
    if len(set(ast.dump(call.args[0]) for call in calls)) != 1:
        raise ValueError("parallel linears read different inputs")
    chains = [block_chain(funcdef.body, call) for call in calls]
    if any(not chain for chain in chains):
        raise ValueError("parallel linear calls not found in the function body")
    blocks = [chain[-1] for chain in chains]
    if len(set(id(block) for block, _ in blocks)) != 1:
        raise ValueError("parallel linear calls are in different blocks")

    block = blocks[0][0]
    indices = [index for _, index in blocks]
    index = min(indices)
    self._check_input_stable(funcdef, module, calls, block, indices)

    # l1(x), l2(x), ... -> merged(x).split(merged.output_sizes / merged.tp_size, -1)
    self._splice_merged_split(funcdef, calls, block, index)
    self.fused_forward = compile_forward(funcdef, fn)

validate(module, vllm_config)

Check that the parallel linears are compatible for merging.

Source code in vllm/model_executor/models/transformers/fusers/merged_column.py 
113
114
115
116
117
118
119
120
121
122
def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
    """Check that the parallel linears are compatible for merging."""
    linear_layers = [module.get_submodule(name) for name in self.linear_names]
    tp_size = vllm_config.parallel_config.tensor_parallel_size
    return (
        len(linear_layers) >= 2
        and len(set(linear.in_features for linear in linear_layers)) == 1
        and len(set(linear.bias is None for linear in linear_layers)) == 1
        and all(linear.out_features % tp_size == 0 for linear in linear_layers)
    )

MoEBlockFuser dataclass

Fuser for MoE block experts, gate and shared_experts (optional).

Methods:

  • gate –

    Rebuild the HF gate as a GateLinear for vLLM's fused MoE.

  • rewrite_forward –

    Rewrite moe_block.forward to route through vLLM's fused MoE.

  • shared_experts –

    Build the HF shared expert (and its optional gate)

Source code in vllm/model_executor/models/transformers/fusers/moe.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
@dataclass
class MoEBlockFuser:
    """Fuser for MoE block `experts`, `gate` and `shared_experts` (optional)."""

    gate_name: str
    scoring_func: str
    shared_name: str | None
    shared_gate_name: str | None
    router_dtype: torch.dtype | None = None

    @staticmethod
    def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None:
        """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`.

        Returns the scoring function and the dtype the router computes in."""
        state = {name for name, _ in named_state(gate)}
        if "weight" not in state or state - {"weight", "e_score_correction_bias"}:
            return None
        graph = trace(gate)
        if graph is None:
            return None
        # The routing top-k is the last one; any earlier one scores expert groups.
        topks = [node for node in graph.nodes if is_op(node, "topk")]
        if not topks:
            return None
        topk = topks[-1]
        # 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]
        logits_cone = _reaches(scorer, "all_input_nodes")
        if not any(is_op(n, "linear") for n in logits_cone):
            return None
        scoring_func = "softmax" if is_op(scorer, "softmax") else "sigmoid"
        return scoring_func, _forced_dtype(logits_cone)

    @staticmethod
    def _match_shared_experts(
        graph: fx.Graph, experts: str
    ) -> tuple[str | None, str | None]:
        """Detects the shared expert and its optional gate by dataflow."""
        experts_predicate = lambda n: n.op == "call_module" and n.target == experts
        if (experts_node := find_node(graph, experts_predicate)) is None:
            return None, None
        from_experts = _reaches(experts_node, "users")
        for add in graph.nodes:
            if not is_op(add, "add"):
                continue
            operands = [a for a in add.args if isinstance(a, fx.Node)]
            # Exactly one side is the experts' output; the other is the shared path.
            sides = [a in from_experts for a in operands]
            if len(operands) != 2 or sides.count(True) != 1:
                continue
            cone = _reaches(operands[sides.index(False)], "all_input_nodes")
            modules = [n for n in cone if n.op == "call_module" and n.target != experts]
            # A sigmoid wrapping one of those modules marks the shared-expert gate.
            gate = next(
                (
                    src
                    for n in cone
                    if is_op(n, "sigmoid")
                    and isinstance(src := peel(n.args[0]), fx.Node)
                    and src in modules
                ),
                None,
            )
            shared = [n for n in modules if n is not gate]
            if len(shared) != 1:
                return None, None
            return shared[0].target, (gate.target if gate is not None else None)
        return None, None

    @classmethod
    def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None":
        # Standard MoE block returns a single tensor.
        if _returns_tuple(type(moe_block)):
            return None
        # Router: the child that scores + top-k selects.
        gate_name = scoring_func = router_dtype = None
        for name, child in moe_block.named_children():
            if (
                name != experts_name
                and (router := cls._match_router(child)) is not None
            ):
                gate_name = name
                scoring_func, router_dtype = router
                break
        if gate_name is None or scoring_func is None:
            return None
        # Shared expert: a child the block adds to the experts' output.
        shared_name = shared_gate_name = None
        others = [
            n
            for n, _ in moe_block.named_children()
            if n not in {experts_name, gate_name}
        ]
        if others:
            graph = trace(moe_block)
            if graph is None:
                return None
            shared_name, shared_gate_name = cls._match_shared_experts(
                graph, experts_name
            )
            if shared_gate_name is not None and not _is_scalar_gate(
                getattr(moe_block, shared_gate_name)
            ):
                return None
        # Fail closed: `rewrite_forward` runs only the experts and the detected
        # shared expert, so any other stateful child would be dropped.
        accounted = {experts_name, gate_name, shared_name, shared_gate_name}
        for name, child in moe_block.named_children():
            if name not in accounted and next(named_state(child), None) is not None:
                return None
        return cls(gate_name, scoring_func, shared_name, shared_gate_name, router_dtype)

    def gate(
        self, moe_block: nn.Module, prefix: str, out_dtype: torch.dtype | None = None
    ) -> GateLinear:
        """Rebuild the HF gate as a `GateLinear` for vLLM's fused MoE."""
        hf_gate = getattr(moe_block, self.gate_name)
        num_experts, hidden_size = hf_gate.weight.shape
        gate = GateLinear(
            hidden_size,
            num_experts,
            bias=False,
            out_dtype=out_dtype or self.router_dtype,
            prefix=maybe_prefix(prefix, self.gate_name),
        )
        if (bias := getattr(hf_gate, "e_score_correction_bias", None)) is not None:
            gate.register_buffer("e_score_correction_bias", bias)
        setattr(moe_block, self.gate_name, gate)
        return gate

    def shared_experts(
        self, moe_block: nn.Module, prefix: str
    ) -> SharedExpertMLP | None:
        """Build the HF shared expert (and its optional gate)
        as a `SharedExpertMLP` for vLLM's fused MoE."""
        if self.shared_name is None:
            return None
        shared_experts = getattr(moe_block, self.shared_name)
        gate = None
        if self.shared_gate_name is not None:
            hf_gate = getattr(moe_block, self.shared_gate_name)
            gate = ReplicatedLinear(
                hf_gate.in_features,
                hf_gate.out_features,
                bias=hf_gate.bias is not None,
                prefix=maybe_prefix(prefix, self.shared_gate_name),
            )
            setattr(moe_block, self.shared_gate_name, gate)
        return SharedExpertMLP(shared_experts, gate)

    def rewrite_forward(self, moe_block: nn.Module) -> None:
        """Rewrite `moe_block.forward` to route through vLLM's fused MoE."""
        moe_block.forward = types.MethodType(_moe_block_forward, moe_block)

_match_router(gate) staticmethod

Matches topk(score(linear(x))), score being softmax/sigmoid.

Returns the scoring function and the dtype the router computes in.

Source code in vllm/model_executor/models/transformers/fusers/moe.py 
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
@staticmethod
def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None:
    """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`.

    Returns the scoring function and the dtype the router computes in."""
    state = {name for name, _ in named_state(gate)}
    if "weight" not in state or state - {"weight", "e_score_correction_bias"}:
        return None
    graph = trace(gate)
    if graph is None:
        return None
    # The routing top-k is the last one; any earlier one scores expert groups.
    topks = [node for node in graph.nodes if is_op(node, "topk")]
    if not topks:
        return None
    topk = topks[-1]
    # 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]
    logits_cone = _reaches(scorer, "all_input_nodes")
    if not any(is_op(n, "linear") for n in logits_cone):
        return None
    scoring_func = "softmax" if is_op(scorer, "softmax") else "sigmoid"
    return scoring_func, _forced_dtype(logits_cone)

_match_shared_experts(graph, experts) staticmethod

Detects the shared expert and its optional gate by dataflow.

Source code in vllm/model_executor/models/transformers/fusers/moe.py 
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
@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

gate(moe_block, prefix, out_dtype=None)

Rebuild the HF gate as a GateLinear for vLLM's fused MoE.

Source code in vllm/model_executor/models/transformers/fusers/moe.py 
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
def gate(
    self, moe_block: nn.Module, prefix: str, out_dtype: torch.dtype | None = None
) -> GateLinear:
    """Rebuild the HF gate as a `GateLinear` for vLLM's fused MoE."""
    hf_gate = getattr(moe_block, self.gate_name)
    num_experts, hidden_size = hf_gate.weight.shape
    gate = GateLinear(
        hidden_size,
        num_experts,
        bias=False,
        out_dtype=out_dtype or self.router_dtype,
        prefix=maybe_prefix(prefix, self.gate_name),
    )
    if (bias := getattr(hf_gate, "e_score_correction_bias", None)) is not None:
        gate.register_buffer("e_score_correction_bias", bias)
    setattr(moe_block, self.gate_name, gate)
    return gate

rewrite_forward(moe_block)

Rewrite moe_block.forward to route through vLLM's fused MoE.

Source code in vllm/model_executor/models/transformers/fusers/moe.py 
299
300
301
def rewrite_forward(self, moe_block: nn.Module) -> None:
    """Rewrite `moe_block.forward` to route through vLLM's fused MoE."""
    moe_block.forward = types.MethodType(_moe_block_forward, moe_block)

shared_experts(moe_block, prefix)

Build the HF shared expert (and its optional gate) as a SharedExpertMLP for vLLM's fused MoE.

Source code in vllm/model_executor/models/transformers/fusers/moe.py 
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
def shared_experts(
    self, moe_block: nn.Module, prefix: str
) -> SharedExpertMLP | None:
    """Build the HF shared expert (and its optional gate)
    as a `SharedExpertMLP` for vLLM's fused MoE."""
    if self.shared_name is None:
        return None
    shared_experts = getattr(moe_block, self.shared_name)
    gate = None
    if self.shared_gate_name is not None:
        hf_gate = getattr(moe_block, self.shared_gate_name)
        gate = ReplicatedLinear(
            hf_gate.in_features,
            hf_gate.out_features,
            bias=hf_gate.bias is not None,
            prefix=maybe_prefix(prefix, self.shared_gate_name),
        )
        setattr(moe_block, self.shared_gate_name, gate)
    return SharedExpertMLP(shared_experts, gate)

PackedQKVFuser dataclass

Bases: RewriteFuser

Fuser for attention with q, k and v packed into one projection.

Methods:

  • update_forward –

    Rewrite the split sizes to the sharded projection's per-rank widths.

  • validate –

    Shapes must be compatible with a head-sharded packed GEMM.

Source code in vllm/model_executor/models/transformers/fusers/packed_qkv.py 
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
@dataclass
class PackedQKVFuser(RewriteFuser):
    """Fuser for attention with q, k and v packed into one projection."""

    qkv_name: str
    o_name: str | None
    q_size: int
    kv_size: int

    def info(self, name: str) -> str:
        return (
            f"Fused: {self.qkv_name} ({name}: {self.source_cls}) -> QKVParallelLinear"
        )

    @staticmethod
    def _packed_sizes(node: fx.Node) -> tuple[int, int] | None:
        """`(q, kv)` from a `split((q, kv, kv), ...)` call, if it is one."""
        if not is_method(node, "split") or len(node.args) < 2:
            return None
        sizes = node.args[1]
        if not isinstance(sizes, (tuple, list)) or len(sizes) != 3:
            return None
        if not all(isinstance(size, int) for size in sizes):
            return None
        q_size, k_size, v_size = sizes
        if k_size != v_size or q_size < k_size:
            return None
        return q_size, k_size

    @classmethod
    def match(cls, graph: fx.Graph, module: nn.Module) -> "PackedQKVFuser | None":
        for node in graph.nodes:
            if (sizes := cls._packed_sizes(node)) is None:
                continue
            q_size, kv_size = sizes
            qkv_node = upstream_linear(node.args[0], module)
            if qkv_node is None:
                continue
            qkv_name = str(qkv_node.target)
            # The split must consume the whole projection.
            if module.get_submodule(qkv_name).out_features != q_size + 2 * kv_size:
                continue
            # o_proj produces the module's output and consumes the query width.
            o_name = returned_linear(graph, module)
            if o_name == qkv_name or (
                o_name is not None
                and module.get_submodule(o_name).in_features != q_size
            ):
                o_name = None
            return cls(
                source_cls=type(module).__name__,
                qkv_name=qkv_name,
                o_name=o_name,
                q_size=q_size,
                kv_size=kv_size,
            )
        return None

    def _split_call(self, funcdef: ast.FunctionDef) -> ast.Call:
        """The unique `self.<qkv_name>(...)....split((a, b, c), ...)` call."""
        calls = [
            node
            for node in ast.walk(funcdef)
            if isinstance(node, ast.Call)
            and isinstance(node.func, ast.Attribute)
            and node.func.attr == "split"
            and node.args
            and isinstance(node.args[0], (ast.Tuple, ast.List))
            and len(node.args[0].elts) == 3
            and any(
                isinstance(inner, ast.Attribute) and inner.attr == self.qkv_name
                for inner in ast.walk(node.func.value)
            )
        ]
        if len(calls) != 1:
            raise ValueError(f"{self.qkv_name} has {len(calls)} three-way splits")
        return calls[0]

    def update_forward(self, module: nn.Module) -> None:
        """Rewrite the split sizes to the sharded projection's per-rank widths."""
        funcdef, fn = recover_forward(type(module))
        split = self._split_call(funcdef)
        # (q, kv, kv) -> [s // qkv.tp_size for s in qkv.output_sizes]
        sections = local_output_sizes(self.qkv_name)
        split.args[0] = ast.parse(sections, mode="eval").body
        self.fused_forward = compile_forward(funcdef, fn)

    def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
        """Shapes must be compatible with a head-sharded packed GEMM."""
        head_size = fused_head_size(module, vllm_config)
        qkv = module.get_submodule(self.qkv_name)
        compatible = (
            self.q_size % head_size == 0
            and self.kv_size % head_size == 0
            and qkv.out_features == self.q_size + 2 * self.kv_size
        )
        if not compatible:
            logger.debug("%s is not compatible with packed QKV fusion", type(module))
        return compatible

    def update_attrs(
        self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
    ) -> None:
        quant_config = vllm_config.quant_config
        head_size = fused_head_size(module, vllm_config)
        qkv_prefix = maybe_prefix(prefix, self.qkv_name)
        qkv = module.get_submodule(self.qkv_name)
        merged = QKVParallelLinear(
            hidden_size=qkv.in_features,
            head_size=head_size,
            total_num_heads=self.q_size // head_size,
            total_num_kv_heads=self.kv_size // head_size,
            bias=qkv.bias is not None,
            quant_config=quant_config,
            prefix=qkv_prefix,
            return_bias=False,
        )
        setattr(module, self.qkv_name, merged)
        log_replacement(qkv_prefix, qkv, merged)
        # If there is an output projection, we know it must be rowwise.
        if self.o_name is not None:
            o_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_prefix
            )
            setattr(module, self.o_name, new_o)
            log_replacement(o_prefix, o_proj, new_o)

_packed_sizes(node) staticmethod

(q, kv) from a split((q, kv, kv), ...) call, if it is one.

Source code in vllm/model_executor/models/transformers/fusers/packed_qkv.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
@staticmethod
def _packed_sizes(node: fx.Node) -> tuple[int, int] | None:
    """`(q, kv)` from a `split((q, kv, kv), ...)` call, if it is one."""
    if not is_method(node, "split") or len(node.args) < 2:
        return None
    sizes = node.args[1]
    if not isinstance(sizes, (tuple, list)) or len(sizes) != 3:
        return None
    if not all(isinstance(size, int) for size in sizes):
        return None
    q_size, k_size, v_size = sizes
    if k_size != v_size or q_size < k_size:
        return None
    return q_size, k_size

_split_call(funcdef)

The unique self.<qkv_name>(...)....split((a, b, c), ...) call.

Source code in vllm/model_executor/models/transformers/fusers/packed_qkv.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def _split_call(self, funcdef: ast.FunctionDef) -> ast.Call:
    """The unique `self.<qkv_name>(...)....split((a, b, c), ...)` call."""
    calls = [
        node
        for node in ast.walk(funcdef)
        if isinstance(node, ast.Call)
        and isinstance(node.func, ast.Attribute)
        and node.func.attr == "split"
        and node.args
        and isinstance(node.args[0], (ast.Tuple, ast.List))
        and len(node.args[0].elts) == 3
        and any(
            isinstance(inner, ast.Attribute) and inner.attr == self.qkv_name
            for inner in ast.walk(node.func.value)
        )
    ]
    if len(calls) != 1:
        raise ValueError(f"{self.qkv_name} has {len(calls)} three-way splits")
    return calls[0]

update_forward(module)

Rewrite the split sizes to the sharded projection's per-rank widths.

Source code in vllm/model_executor/models/transformers/fusers/packed_qkv.py 
115
116
117
118
119
120
121
122
def update_forward(self, module: nn.Module) -> None:
    """Rewrite the split sizes to the sharded projection's per-rank widths."""
    funcdef, fn = recover_forward(type(module))
    split = self._split_call(funcdef)
    # (q, kv, kv) -> [s // qkv.tp_size for s in qkv.output_sizes]
    sections = local_output_sizes(self.qkv_name)
    split.args[0] = ast.parse(sections, mode="eval").body
    self.fused_forward = compile_forward(funcdef, fn)

validate(module, vllm_config)

Shapes must be compatible with a head-sharded packed GEMM.

Source code in vllm/model_executor/models/transformers/fusers/packed_qkv.py 
124
125
126
127
128
129
130
131
132
133
134
135
def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
    """Shapes must be compatible with a head-sharded packed GEMM."""
    head_size = fused_head_size(module, vllm_config)
    qkv = module.get_submodule(self.qkv_name)
    compatible = (
        self.q_size % head_size == 0
        and self.kv_size % head_size == 0
        and qkv.out_features == self.q_size + 2 * self.kv_size
    )
    if not compatible:
        logger.debug("%s is not compatible with packed QKV fusion", type(module))
    return compatible

QKVFuser dataclass

Bases: StackedFuser

Fuser for the attention QKV pattern q(x), k(x), v(x).

Methods:

  • update_forward –

    Replace q(x), k(x), v(x) with qkv(x).split(sizes, -1) in source.

  • validate –

    Shapes must be compatible for a single merged, head-sharded GEMM.

Source code in vllm/model_executor/models/transformers/fusers/qkv.py 
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
@dataclass
class QKVFuser(StackedFuser):
    """Fuser for the attention QKV pattern `q(x), k(x), v(x)`."""

    q_name: str
    k_name: str
    v_name: str
    o_name: str | None
    merged_name: ClassVar[str] = "qkv_proj"
    merged_cls_name: ClassVar[str] = "QKVParallelLinear"

    @property
    def shards(self) -> list[tuple[str, ShardId]]:
        return [(self.q_name, "q"), (self.k_name, "k"), (self.v_name, "v")]

    @classmethod
    def _get_qkv_nodes(
        cls, graph: fx.Graph, module: nn.Module
    ) -> tuple[fx.Node, fx.Node, fx.Node] | None:
        """Search `graph` for the QKV pattern `q(x), k(x), v(x)`."""
        by_input: dict[fx.Node, list[fx.Node]] = {}
        for node in graph.nodes:
            if (
                is_linear(node, module)
                and len(node.args) == 1
                and not node.kwargs
                and isinstance(node.args[0], fx.Node)
                and node.args[0].op == "placeholder"
            ):
                by_input.setdefault(node.args[0], []).append(node)
        triples = [nodes for nodes in by_input.values() if len(nodes) == 3]
        if len(triples) != 1:
            return None

        q_node, k_node, v_node = nodes = triples[0]
        outs = [module.get_submodule(node.target).out_features for node in nodes]
        if len(set(outs)) == 2:
            # q is identified as the larger projection (GQA)
            (q_node,) = (n for n, out in zip(nodes, outs) if outs.count(out) == 1)
            k_node, v_node = (n for n, out in zip(nodes, outs) if outs.count(out) == 2)
            if module.get_submodule(q_node.target).out_features != max(outs):
                return None
        elif len(set(outs)) != 1:
            return None
        return q_node, k_node, v_node

    @classmethod
    def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None":
        if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None:
            return None
        q, k, v = qkv_nodes
        names = dict(q_name=q.target, k_name=k.target, v_name=v.target)
        # o_proj produces the module's output.
        o_name = returned_linear(graph, module)
        # o_proj must be compatible with the q/k/v projections.
        if o_name in names.values() or (
            o_name is not None
            and module.get_submodule(o_name).in_features
            != module.get_submodule(q.target).out_features
        ):
            o_name = None
        names["o_name"] = o_name
        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.

        A projection may be guarded by an existence check -- a
        `self.<proj> is (not) None` comparison (e.g. Gemma 4's `v_proj`) or a bare
        truthiness test -- which is folded to its constant value (see
        `bypass_existence_guard`). The calls may sit in different branches, so the
        fused GEMM is inserted before the earliest of them, in the innermost block
        that dominates all three.
        """
        funcdef, fn = recover_forward(type(module))
        proj_names = (self.q_name, self.k_name, self.v_name)
        calls = self._unguarded_calls(funcdef, proj_names)
        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)
        } - set(proj_names)
        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")

        # Insert the fused GEMM before the earliest call, in the innermost block
        # common to all three (the calls may be split across branches).
        chains = [block_chain(funcdef.body, call) for call in calls]
        if any(not chain for chain in chains):
            raise ValueError("projection calls not found in the function body")
        depth = 0
        for level in zip(*chains):
            if len({id(block) for block, _ in level}) != 1:
                break
            depth += 1
        block = chains[0][depth - 1][0]
        indices = [chain[depth - 1][1] for chain in chains]
        insert_index = min(indices)
        self._check_input_stable(funcdef, module, calls, block, indices)

        # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
        self._splice_merged_split(funcdef, calls, block, insert_index)
        self.fused_forward = compile_forward(funcdef, fn)

    def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> 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 = fused_head_size(module, vllm_config)
        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, vllm_config: "VllmConfig"
    ) -> None:
        quant_config = vllm_config.quant_config
        head_size = fused_head_size(module, vllm_config)
        q = module.get_submodule(self.q_name)
        k = module.get_submodule(self.k_name)
        merged = QKVParallelLinear(
            hidden_size=q.in_features,
            head_size=head_size,
            total_num_heads=q.out_features // head_size,
            total_num_kv_heads=k.out_features // head_size,
            bias=q.bias is not None,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, self.merged_name),
            return_bias=False,
        )
        logger.debug(
            "%s: %s, %s: %s, %s: %s -> %s: %s",
            self.q_name,
            q,
            self.k_name,
            k,
            self.v_name,
            module.get_submodule(self.v_name),
            self.merged_name,
            merged,
        )
        setattr(module, self.merged_name, merged)
        # Drop the consumed submodules so their (meta) params are not expected.
        for name in (self.q_name, self.k_name, self.v_name):
            delattr(module, name)
        # If there is an output projection, we know it must be rowwise.
        if self.o_name is not None:
            o_proj_prefix = maybe_prefix(prefix, self.o_name)
            o_proj = module.get_submodule(self.o_name)
            new_o = replace_linear_class(
                o_proj, "rowwise", quant_config, prefix=o_proj_prefix
            )
            setattr(module, self.o_name, new_o)
            log_replacement(o_proj_prefix, o_proj, new_o)

_get_qkv_nodes(graph, module) classmethod

Search graph for the QKV pattern q(x), k(x), v(x).

Source code in vllm/model_executor/models/transformers/fusers/qkv.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
@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

update_forward(module)

Replace q(x), k(x), v(x) with qkv(x).split(sizes, -1) in source.

A projection may be guarded by an existence check -- a self.<proj> is (not) None comparison (e.g. Gemma 4's v_proj) or a bare truthiness test -- which is folded to its constant value (see bypass_existence_guard). The calls may sit in different branches, so the fused GEMM is inserted before the earliest of them, in the innermost block that dominates all three.

Source code in vllm/model_executor/models/transformers/fusers/qkv.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def update_forward(self, module: nn.Module) -> None:
    """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source.

    A projection may be guarded by an existence check -- a
    `self.<proj> is (not) None` comparison (e.g. Gemma 4's `v_proj`) or a bare
    truthiness test -- which is folded to its constant value (see
    `bypass_existence_guard`). The calls may sit in different branches, so the
    fused GEMM is inserted before the earliest of them, in the innermost block
    that dominates all three.
    """
    funcdef, fn = recover_forward(type(module))
    proj_names = (self.q_name, self.k_name, self.v_name)
    calls = self._unguarded_calls(funcdef, proj_names)
    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)
    } - set(proj_names)
    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")

    # Insert the fused GEMM before the earliest call, in the innermost block
    # common to all three (the calls may be split across branches).
    chains = [block_chain(funcdef.body, call) for call in calls]
    if any(not chain for chain in chains):
        raise ValueError("projection calls not found in the function body")
    depth = 0
    for level in zip(*chains):
        if len({id(block) for block, _ in level}) != 1:
            break
        depth += 1
    block = chains[0][depth - 1][0]
    indices = [chain[depth - 1][1] for chain in chains]
    insert_index = min(indices)
    self._check_input_stable(funcdef, module, calls, block, indices)

    # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
    self._splice_merged_split(funcdef, calls, block, insert_index)
    self.fused_forward = compile_forward(funcdef, fn)

validate(module, vllm_config)

Shapes must be compatible for a single merged, head-sharded GEMM.

Source code in vllm/model_executor/models/transformers/fusers/qkv.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> 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 = fused_head_size(module, vllm_config)
    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

RMSNormFuser dataclass

Bases: BaseFuser

Fuser for RMSNorm patterns, including Gemma-style zero-centered weights.

Methods:

  • fuse –

    Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp.

  • match –

    Match a graph to the RMSNorm pattern, returning a fuser if found.

Attributes:

  • eps (float | None) –

    eps itself, when it is not held in an attribute (see _eps_source).

  • eps_attr (str | None) –

    Attribute holding eps, read per instance in fuse.

  • source_cls (str) –

    Class name of the norm this was matched from (for logging).

  • zero_centered (bool) –

    Gemma-style (1 + weight) scaling (weight initialised at zero).

Source code in vllm/model_executor/models/transformers/fusers/rms_norm.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
@dataclass
class RMSNormFuser(BaseFuser):
    """Fuser for RMSNorm patterns, including Gemma-style zero-centered weights."""

    zero_centered: bool
    """Gemma-style `(1 + weight)` scaling (weight initialised at zero)."""
    source_cls: str
    """Class name of the norm this was matched from (for logging)."""
    eps_attr: str | None = None
    """Attribute holding `eps`, read per instance in `fuse`."""
    eps: float | None = None
    """`eps` itself, when it is not held in an attribute (see `_eps_source`)."""

    def info(self, name: str) -> str:
        norm = "GemmaRMSNorm" if self.zero_centered else "RMSNorm"
        return f"Fused: {name} ({self.source_cls}) -> {norm} (CustomOp)"

    @classmethod
    def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
        """Match a graph to the RMSNorm pattern, returning a fuser if found."""
        if forward_input_count(type(module)) != 1:
            return None
        x = find_node(graph, lambda n: n.op == "placeholder")
        if x is None:
            return None
        # Handle native torch `rms_norm` op.
        rms_norm = find_node(graph, lambda n: is_op(n, "rms_norm"))
        if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x:
            if _has_trailing_compute(graph, rms_norm):
                return None
            eps_attr, eps = cls._eps_source(graph, module)
            return cls(
                zero_centered=False,
                source_cls=type(module).__name__,
                eps_attr=eps_attr,
                eps=eps,
            )
        # Handle explicit `x * rsqrt(mean(x**2, -1) + eps)` pattern.
        # The rsqrt over the mean-square variance is the spine of the norm.
        rsqrt = None
        for node in graph.nodes:
            if is_op(node, "rsqrt") and _variance_eps(node, x) is not None:
                rsqrt = node
                break
        if rsqrt is None:
            return None
        # The `x * rsqrt(...)` normalize multiply.
        normalize = find_node(
            graph, lambda n: is_op(n, "mul") and rsqrt in map(peel, n.args)
        )
        if normalize is None:
            return None
        # An optional later `weight * normalized` (or `(1 + weight) * normalized`).
        tail, zero_centered = normalize, False
        for node in graph.nodes:
            if not is_op(node, "mul") or node is normalize:
                continue
            operands = [peel(a) for a in node.args if isinstance(a, fx.Node)]
            if len(operands) == 2 and normalize in operands:
                weight = next(o for o in operands if o is not normalize)
                tail, zero_centered = node, _is_one_plus(weight)
                break
        # The norm must be the last compute in forward, or it is not a pure norm.
        if _has_trailing_compute(graph, tail):
            return None
        eps_attr, eps = cls._eps_source(graph, module)
        return cls(
            zero_centered=zero_centered,
            source_cls=type(module).__name__,
            eps_attr=eps_attr,
            eps=eps,
        )

    @classmethod
    def _eps_source(
        cls, graph: fx.Graph, module: nn.Module
    ) -> tuple[str | None, float | None]:
        """Where `fuse` should read `eps` from, resolved once per class."""
        eps = cls._eps_from_graph(graph)
        if eps is None:
            return None, None
        # Whatever supplied the constant must still equal it.
        candidates = {
            name: value
            for name, value in vars(module).items()
            if isinstance(value, float) and value == eps
        }
        # Use unique markers and retrace to verify exactly which attribute is eps.
        markers = {float(-index - 1): name for index, name in enumerate(candidates)}
        marked = None
        if markers:
            try:
                for marker, name in markers.items():
                    setattr(module, name, marker)
                if (remarked := trace(module)) is not None:
                    marked = cls._eps_from_graph(remarked)
            finally:
                for name, value in candidates.items():
                    setattr(module, name, value)
        if (name := markers.get(marked)) is not None:
            return name, None
        logger.debug_once(
            "%s does not hold its eps (%s) in an attribute. Every instance in this "
            "model will use the value traced from this instance. If this is not "
            "desired, consider storing and reading eps using attribute of %s.",
            type(module).__name__,
            eps,
            type(module).__name__,
        )
        return None, eps

    @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 float(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, vllm_config: "VllmConfig") -> bool:
        return True

    def fuse(
        self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
    ) -> nn.Module:
        """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp."""
        weight = getattr(module, "weight", None)
        has_weight = weight is not None
        hidden_size = weight.size(0) if has_weight else 0
        eps = getattr(module, self.eps_attr, None) if self.eps_attr else self.eps
        if not isinstance(eps, (int, float)):
            # If eps was not detected, match torch behaviour.
            dtype = weight.dtype if has_weight else vllm_config.model_config.dtype
            eps = torch.finfo(dtype).eps
        if self.zero_centered:
            return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps)
        return TPAwareRMSNorm(
            hidden_size=hidden_size,
            eps=eps,
            has_weight=has_weight,
            dtype=weight.dtype if has_weight else None,
        )

eps = None class-attribute instance-attribute

eps itself, when it is not held in an attribute (see _eps_source).

eps_attr = None class-attribute instance-attribute

Attribute holding eps, read per instance in fuse.

source_cls instance-attribute

Class name of the norm this was matched from (for logging).

zero_centered instance-attribute

Gemma-style (1 + weight) scaling (weight initialised at zero).

_eps_from_graph(graph) staticmethod

Extract the eps constant from the graph, if present.

Source code in vllm/model_executor/models/transformers/fusers/rms_norm.py 
231
232
233
234
235
236
237
238
239
240
241
242
243
244
@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 float(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

_eps_source(graph, module) classmethod

Where fuse should read eps from, resolved once per class.

Source code in vllm/model_executor/models/transformers/fusers/rms_norm.py 
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
@classmethod
def _eps_source(
    cls, graph: fx.Graph, module: nn.Module
) -> tuple[str | None, float | None]:
    """Where `fuse` should read `eps` from, resolved once per class."""
    eps = cls._eps_from_graph(graph)
    if eps is None:
        return None, None
    # Whatever supplied the constant must still equal it.
    candidates = {
        name: value
        for name, value in vars(module).items()
        if isinstance(value, float) and value == eps
    }
    # Use unique markers and retrace to verify exactly which attribute is eps.
    markers = {float(-index - 1): name for index, name in enumerate(candidates)}
    marked = None
    if markers:
        try:
            for marker, name in markers.items():
                setattr(module, name, marker)
            if (remarked := trace(module)) is not None:
                marked = cls._eps_from_graph(remarked)
        finally:
            for name, value in candidates.items():
                setattr(module, name, value)
    if (name := markers.get(marked)) is not None:
        return name, None
    logger.debug_once(
        "%s does not hold its eps (%s) in an attribute. Every instance in this "
        "model will use the value traced from this instance. If this is not "
        "desired, consider storing and reading eps using attribute of %s.",
        type(module).__name__,
        eps,
        type(module).__name__,
    )
    return None, eps

fuse(module, prefix, vllm_config)

Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp.

Source code in vllm/model_executor/models/transformers/fusers/rms_norm.py 
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
def fuse(
    self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
) -> nn.Module:
    """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp."""
    weight = getattr(module, "weight", None)
    has_weight = weight is not None
    hidden_size = weight.size(0) if has_weight else 0
    eps = getattr(module, self.eps_attr, None) if self.eps_attr else self.eps
    if not isinstance(eps, (int, float)):
        # If eps was not detected, match torch behaviour.
        dtype = weight.dtype if has_weight else vllm_config.model_config.dtype
        eps = torch.finfo(dtype).eps
    if self.zero_centered:
        return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps)
    return TPAwareRMSNorm(
        hidden_size=hidden_size,
        eps=eps,
        has_weight=has_weight,
        dtype=weight.dtype if has_weight else None,
    )

match(graph, module) classmethod

Match a graph to the RMSNorm pattern, returning a fuser if found.

Source code in vllm/model_executor/models/transformers/fusers/rms_norm.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
@classmethod
def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
    """Match a graph to the RMSNorm pattern, returning a fuser if found."""
    if forward_input_count(type(module)) != 1:
        return None
    x = find_node(graph, lambda n: n.op == "placeholder")
    if x is None:
        return None
    # Handle native torch `rms_norm` op.
    rms_norm = find_node(graph, lambda n: is_op(n, "rms_norm"))
    if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x:
        if _has_trailing_compute(graph, rms_norm):
            return None
        eps_attr, eps = cls._eps_source(graph, module)
        return cls(
            zero_centered=False,
            source_cls=type(module).__name__,
            eps_attr=eps_attr,
            eps=eps,
        )
    # Handle explicit `x * rsqrt(mean(x**2, -1) + eps)` pattern.
    # The rsqrt over the mean-square variance is the spine of the norm.
    rsqrt = None
    for node in graph.nodes:
        if is_op(node, "rsqrt") and _variance_eps(node, x) is not None:
            rsqrt = node
            break
    if rsqrt is None:
        return None
    # The `x * rsqrt(...)` normalize multiply.
    normalize = find_node(
        graph, lambda n: is_op(n, "mul") and rsqrt in map(peel, n.args)
    )
    if normalize is None:
        return None
    # An optional later `weight * normalized` (or `(1 + weight) * normalized`).
    tail, zero_centered = normalize, False
    for node in graph.nodes:
        if not is_op(node, "mul") or node is normalize:
            continue
        operands = [peel(a) for a in node.args if isinstance(a, fx.Node)]
        if len(operands) == 2 and normalize in operands:
            weight = next(o for o in operands if o is not normalize)
            tail, zero_centered = node, _is_one_plus(weight)
            break
    # The norm must be the last compute in forward, or it is not a pure norm.
    if _has_trailing_compute(graph, tail):
        return None
    eps_attr, eps = cls._eps_source(graph, module)
    return cls(
        zero_centered=zero_centered,
        source_cls=type(module).__name__,
        eps_attr=eps_attr,
        eps=eps,
    )

RewriteFuser dataclass

Bases: BaseFuser

A fuser that rewrites the module's forward and rebinds it.

match and update_forward analyse the class once; fuse swaps the submodules and binds the compiled forward on an instance in place, so it keeps its class and any attribute the fusion does not consume.

Methods:

  • fuse –

    Fuse an already-validated module in place (see Fusers.__getitem__).

  • update_attrs –

    Replace module's submodules with their vLLM equivalents.

  • update_forward –

    Rewrite and compile type(module)'s forward source.

Attributes:

  • fused_forward (Callable) –

    The compiled rewritten forward, set by update_forward.

  • source_cls (str) –

    Class of the HF module the fused projections belonged to (for logging).

Source code in vllm/model_executor/models/transformers/fusers/base.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
@dataclass
class RewriteFuser(BaseFuser):
    """A fuser that rewrites the module's forward and rebinds it.

    `match` and `update_forward` analyse the class once; `fuse` swaps the
    submodules and binds the compiled forward on an instance in place, so it
    keeps its class and any attribute the fusion does not consume.
    """

    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`."""

    @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, vllm_config: "VllmConfig"
    ) -> None:
        """Replace `module`'s submodules with their vLLM equivalents."""

    def fuse(
        self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
    ) -> 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, vllm_config)
        module.forward = types.MethodType(self.fused_forward, module)
        return module

fused_forward = field(init=False, repr=False) class-attribute instance-attribute

The compiled rewritten forward, set by update_forward.

source_cls instance-attribute

Class of the HF module the fused projections belonged to (for logging).

fuse(module, prefix, vllm_config)

Fuse an already-validated module in place (see Fusers.__getitem__).

Builds the merged submodule and binds the compiled forward.

Source code in vllm/model_executor/models/transformers/fusers/base.py 
117
118
119
120
121
122
123
124
125
def fuse(
    self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
) -> 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, vllm_config)
    module.forward = types.MethodType(self.fused_forward, module)
    return module

update_attrs(module, prefix, vllm_config) abstractmethod

Replace module's submodules with their vLLM equivalents.

Source code in vllm/model_executor/models/transformers/fusers/base.py 
111
112
113
114
115
@abstractmethod
def update_attrs(
    self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
) -> None:
    """Replace `module`'s submodules with their vLLM equivalents."""

update_forward(module) abstractmethod

Rewrite and compile type(module)'s forward source.

Raises if the source does not admit the rewrite (fusion is then skipped).

Source code in vllm/model_executor/models/transformers/fusers/base.py 
104
105
106
107
108
109
@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).
    """

StackedFuser dataclass

Bases: RewriteFuser

A fuser that merges sibling projections into one stacked linear and rewrites the forward to call it.

Methods:

Attributes:

Source code in vllm/model_executor/models/transformers/fusers/base.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
@dataclass
class StackedFuser(RewriteFuser):
    """A fuser that merges sibling projections into one stacked linear and
    rewrites the forward to call it."""

    merged_name: ClassVar[str]
    """Attribute name of the merged module created by `update_attrs`."""
    merged_cls_name: ClassVar[str]
    """Name of the vLLM class the merged projection becomes (for logging)."""

    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_name})"
        )

    @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]}

    def _unguarded_calls(
        self, funcdef: ast.FunctionDef, names: Iterable[str]
    ) -> list[ast.Call]:
        """One `self.<name>(arg)` call per projection, existence guards folded.

        `update_attrs` deletes the projections it merges, so any reference to one beyond
        its call site must be a guard on its existence; folding those to their constant
        value keeps the rewritten forward off a name that no longer exists. A reference
        that is not such a guard raises, so fusion is skipped."""
        calls = []
        for name in names:
            call, refs = self_call_and_refs(funcdef, name)
            for ref in refs:
                bypass_existence_guard(funcdef, ref, name)
            calls.append(call)
        return calls

    def _check_input_stable(
        self,
        funcdef: ast.FunctionDef,
        module: nn.Module,
        calls: list[ast.Call],
        block: list[ast.stmt],
        indices: list[int],
    ) -> None:
        """Raise unless hoisting the merged GEMM preserves what it reads.

        Fusing moves every projection to one call at `min(indices)`, so the
        merged GEMM reads the input once, up front, where the last of `calls`
        would have read it later. That holds only if nothing in between changes
        the input, and a change need not name it: writing any view that shares
        its storage changes it too. Both halves of the check over-approximate,
        since a false hit costs a fusion while a miss returns wrong numbers.
        """
        arg_names = {
            node.id
            for node in ast.walk(calls[0].args[0])
            if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load)
        }
        # Protect the names the input reads and anything that may alias them.
        tracked = aliasing_names(funcdef, arg_names, module)
        # Any of them rebound, written through, or passed to a call that writes
        # its argument would leave the hoisted GEMM reading a different value.
        region = block[min(indices) : max(indices) + 1]
        if tracked & written_names(region):
            raise ValueError("projection input is rebound or mutated before all calls")

    def _splice_merged_split(
        self,
        funcdef: ast.FunctionDef,
        calls: list[ast.Call],
        block: list[ast.stmt],
        index: int,
    ) -> None:
        """Insert `temps = self.<merged_name>(arg).split(sizes, -1)` at
        `block[index]` and replace each of `calls` with its temp name.

        `calls` must share one input argument; `block[index]` must be where
        they are (or would be) evaluated. Raises if a generated temporary
        would shadow an existing name in `funcdef`.
        """
        temps = [f"_vllm_merged_{i}" for i in range(len(calls))]
        names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
        if names & set(temps):
            raise ValueError("fused temporaries would shadow existing names")
        targets = ", ".join(temps)
        sections = local_output_sizes(self.merged_name)
        source = f"{targets} = self.{self.merged_name}(__arg__).split({sections}, -1)"
        assign = ast.parse(source).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])
        ast.copy_location(assign, block[index])
        block.insert(index, assign)
        for call, temp in zip(calls, temps):
            replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))

merged_cls_name class-attribute

Name of the vLLM class the merged projection becomes (for logging).

merged_name class-attribute

Attribute name of the merged module created by update_attrs.

packed_modules_mapping property

{merged_name: [projection names]} so quantization can unpack the fused layer into its per-shard configs.

shards abstractmethod property

Each projection's original name and its shard id in the merged module.

Source for both orig_to_new_stacked and packed_modules_mapping.

_check_input_stable(funcdef, module, calls, block, indices)

Raise unless hoisting the merged GEMM preserves what it reads.

Fusing moves every projection to one call at min(indices), so the merged GEMM reads the input once, up front, where the last of calls would have read it later. That holds only if nothing in between changes the input, and a change need not name it: writing any view that shares its storage changes it too. Both halves of the check over-approximate, since a false hit costs a fusion while a miss returns wrong numbers.

Source code in vllm/model_executor/models/transformers/fusers/base.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
def _check_input_stable(
    self,
    funcdef: ast.FunctionDef,
    module: nn.Module,
    calls: list[ast.Call],
    block: list[ast.stmt],
    indices: list[int],
) -> None:
    """Raise unless hoisting the merged GEMM preserves what it reads.

    Fusing moves every projection to one call at `min(indices)`, so the
    merged GEMM reads the input once, up front, where the last of `calls`
    would have read it later. That holds only if nothing in between changes
    the input, and a change need not name it: writing any view that shares
    its storage changes it too. Both halves of the check over-approximate,
    since a false hit costs a fusion while a miss returns wrong numbers.
    """
    arg_names = {
        node.id
        for node in ast.walk(calls[0].args[0])
        if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load)
    }
    # Protect the names the input reads and anything that may alias them.
    tracked = aliasing_names(funcdef, arg_names, module)
    # Any of them rebound, written through, or passed to a call that writes
    # its argument would leave the hoisted GEMM reading a different value.
    region = block[min(indices) : max(indices) + 1]
    if tracked & written_names(region):
        raise ValueError("projection input is rebound or mutated before all calls")

_splice_merged_split(funcdef, calls, block, index)

Insert temps = self.<merged_name>(arg).split(sizes, -1) at block[index] and replace each of calls with its temp name.

calls must share one input argument; block[index] must be where they are (or would be) evaluated. Raises if a generated temporary would shadow an existing name in funcdef.

Source code in vllm/model_executor/models/transformers/fusers/base.py 
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
def _splice_merged_split(
    self,
    funcdef: ast.FunctionDef,
    calls: list[ast.Call],
    block: list[ast.stmt],
    index: int,
) -> None:
    """Insert `temps = self.<merged_name>(arg).split(sizes, -1)` at
    `block[index]` and replace each of `calls` with its temp name.

    `calls` must share one input argument; `block[index]` must be where
    they are (or would be) evaluated. Raises if a generated temporary
    would shadow an existing name in `funcdef`.
    """
    temps = [f"_vllm_merged_{i}" for i in range(len(calls))]
    names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
    if names & set(temps):
        raise ValueError("fused temporaries would shadow existing names")
    targets = ", ".join(temps)
    sections = local_output_sizes(self.merged_name)
    source = f"{targets} = self.{self.merged_name}(__arg__).split({sections}, -1)"
    assign = ast.parse(source).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])
    ast.copy_location(assign, block[index])
    block.insert(index, assign)
    for call, temp in zip(calls, temps):
        replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))

_unguarded_calls(funcdef, names)

One self.<name>(arg) call per projection, existence guards folded.

update_attrs deletes the projections it merges, so any reference to one beyond its call site must be a guard on its existence; folding those to their constant value keeps the rewritten forward off a name that no longer exists. A reference that is not such a guard raises, so fusion is skipped.

Source code in vllm/model_executor/models/transformers/fusers/base.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
def _unguarded_calls(
    self, funcdef: ast.FunctionDef, names: Iterable[str]
) -> list[ast.Call]:
    """One `self.<name>(arg)` call per projection, existence guards folded.

    `update_attrs` deletes the projections it merges, so any reference to one beyond
    its call site must be a guard on its existence; folding those to their constant
    value keeps the rewritten forward off a name that no longer exists. A reference
    that is not such a guard raises, so fusion is skipped."""
    calls = []
    for name in names:
        call, refs = self_call_and_refs(funcdef, name)
        for ref in refs:
            bypass_existence_guard(funcdef, ref, name)
        calls.append(call)
    return calls

orig_to_new_stacked(prefix)

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).

Source code in vllm/model_executor/models/transformers/fusers/base.py 
152
153
154
155
156
157
158
159
160
161
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
    }