Skip to content

vllm.model_executor.models.transformers.fusers.merged_column

Fuser for parallel linear projections.

Classes:

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