vllm.model_executor.models.transformers.fusers ¶

@dataclass
class BaseFuser(ABC):
    """A detected fusion and how to apply it.

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

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

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

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

    @abstractmethod
    def fuse(
        self,
        module: nn.Module,
        prefix: str,
        model_config: "ModelConfig",
        quant_config: "QuantizationConfig",
    ) -> nn.Module:
        """Apply the fusion to an already-validated `module`, returning the
        module to install in its place (mutated in place, or freshly built)."""

    def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
        """`WeightsMapper.orig_to_new_stacked` entries this fuser contributes
        (none unless it stacks weights)."""
        return {}

    @property
    def packed_modules_mapping(self) -> dict[str, list[str]]:
        """`packed_modules_mapping` entries this fuser contributes (none unless
        it stacks weights)."""
        return {}

@dataclass
class GLUFuser(StackedFuser):
    """Fuser for the GLU pattern `act(gate(x)) * up(x)`."""

    act_name: str
    gate_name: str
    up_name: str
    down_name: str | None
    merged_name: ClassVar[str] = "gate_up_proj"
    merged_cls: ClassVar[str] = "MergedColumnParallelLinear"

    @property
    def shards(self) -> list[tuple[str, ShardId]]:
        return [(self.gate_name, 0), (self.up_name, 1)]

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

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

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

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

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

        gate = module.get_submodule(gate_node.target)
        up = module.get_submodule(up_node.target)
        # Shapes must be compatible for a single merged GEMM.
        if gate.in_features == up.in_features and (gate.bias is None) == (
            up.bias is None
        ):
            predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node
            down_node = find_node(graph, predicate)
            return cls(
                source_cls=type(module).__name__,
                act_name=act_node.target,
                gate_name=gate_node.target,
                up_name=up_node.target,
                down_name=down_node.target if down_node is not None else None,
            )
        return None

    def update_forward(self, module: nn.Module) -> None:
        """Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source."""
        funcdef, fn = recover_forward(type(module))
        act_call = single_self_call(funcdef, self.act_name)
        gate_call = single_self_call(funcdef, self.gate_name)
        up_call = single_self_call(funcdef, self.up_name)
        if act_call.args[0] is not gate_call:
            raise ValueError("activation does not directly wrap the gate")
        if ast.dump(gate_call.args[0]) != ast.dump(up_call.args[0]):
            raise ValueError("gate and up inputs are written differently")
        muls = [
            node
            for node in ast.walk(funcdef)
            if isinstance(node, ast.BinOp)
            and isinstance(node.op, ast.Mult)
            and {id(node.left), id(node.right)} == {id(act_call), id(up_call)}
        ]
        if len(muls) != 1:
            raise ValueError("no multiply of the activation and up projection")

        # act(gate(x)) * up(x) -> act(gate_up(x))
        assert isinstance(gate_call.func, ast.Attribute)
        gate_call.func.attr = self.merged_name
        replace_expr(funcdef, muls[0], act_call)
        self.fused_forward = compile_forward(funcdef, fn)

    def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
        act = module.get_submodule(self.act_name)
        if self._get_act_and_mul_name(act) is None:
            logger.debug("No AndMul equivalent for %s; skipping fusion", type(act))
            return False
        return True

    def update_attrs(
        self,
        module: nn.Module,
        prefix: str,
        model_config: "ModelConfig",
        quant_config: "QuantizationConfig",
    ) -> None:
        act_fn = self._get_act_and_mul(module.get_submodule(self.act_name))
        gate = module.get_submodule(self.gate_name)
        up = module.get_submodule(self.up_name)
        merged = MergedColumnParallelLinear(
            input_size=gate.in_features,
            output_sizes=[gate.out_features, up.out_features],
            bias=gate.bias is not None,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, self.merged_name),
            return_bias=False,
        )
        logger.debug(
            "%s: %s, %s: %s -> %s: %s",
            self.gate_name,
            gate,
            self.up_name,
            up,
            self.merged_name,
            merged,
        )
        setattr(module, self.merged_name, merged)
        setattr(module, self.act_name, act_fn)
        # Drop the consumed submodules so their (meta) params are not expected.
        delattr(module, self.gate_name)
        delattr(module, self.up_name)
        # If there is a down projection, we know it must be rowwise.
        if self.down_name is not None:
            down_prefix = maybe_prefix(prefix, self.down_name)
            down = module.get_submodule(self.down_name)
            new_down = replace_linear_class(
                down, "rowwise", quant_config, prefix=down_prefix
            )
            setattr(module, self.down_name, new_down)
            log_replacement(down_prefix, down, new_down)

@dataclass
class MoEBlockFuser:
    """Fuser for MoE block `experts`, `gate` and `shared_experts` (optional)."""

    gate_name: str
    scoring_func: str
    shared_name: str | None
    shared_gate_name: str | None

    @staticmethod
    def _match_router(gate: nn.Module) -> str | None:
        """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`."""
        if [name for name, _ in named_state(gate)] != ["weight"]:
            return None
        graph = trace(gate)
        if graph is None:
            return None
        topk = find_node(graph, lambda n: is_op(n, "topk"))
        if topk is None:
            return None
        # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear.
        scorers = [
            n
            for n in _reaches(topk, "all_input_nodes")
            if is_op(n, "softmax") or is_op(n, "sigmoid")
        ]
        if len(scorers) != 1:
            return None
        scorer = scorers[0]
        if not any(is_op(n, "linear") for n in _reaches(scorer, "all_input_nodes")):
            return None
        return "softmax" if is_op(scorer, "softmax") else "sigmoid"

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

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

    def gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear:
        """Rebuild the HF gate as a `ReplicatedLinear` for vLLM's fused MoE."""
        num_experts, hidden_size = getattr(moe_block, self.gate_name).weight.shape
        gate = ReplicatedLinear(
            hidden_size,
            num_experts,
            bias=False,
            prefix=maybe_prefix(prefix, self.gate_name),
        )
        setattr(moe_block, self.gate_name, gate)
        return gate

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

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

@dataclass
class QKVFuser(StackedFuser):
    """Fuser for the attention QKV pattern `q(x), k(x), v(x)`."""

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

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

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

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

    @classmethod
    def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None":
        if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None:
            return None
        q, k, v = qkv_nodes
        names = dict(q_name=q.target, k_name=k.target, v_name=v.target)
        attn_width = module.get_submodule(q.target).out_features
        candidates = [
            name
            for name, child in module.named_children()
            if isinstance(child, nn.Linear)
            and name not in names.values()
            and child.in_features == attn_width
        ]
        names["o_name"] = candidates[0] if len(candidates) == 1 else None
        return cls(source_cls=type(module).__name__, **names)

    def update_forward(self, module: nn.Module) -> None:
        """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source."""
        funcdef, fn = recover_forward(type(module))
        calls = [
            single_self_call(funcdef, name)
            for name in (self.q_name, self.k_name, self.v_name)
        ]
        arg_dumps = {ast.dump(call.args[0]) for call in calls}
        if len(arg_dumps) != 1:
            raise ValueError("projection inputs are written differently")
        # The trace may be partial, so prove projection exclusivity in source:
        # no other linear child may consume the same input (else the matched
        # three may not be q, k and v)
        other_linears = {
            name
            for name, child in module.named_children()
            if isinstance(child, nn.Linear)
        } - {self.q_name, self.k_name, self.v_name}
        for node in ast.walk(funcdef):
            if (
                isinstance(node, ast.Call)
                and isinstance(node.func, ast.Attribute)
                and node.func.attr in other_linears
                and any(ast.dump(arg) in arg_dumps for arg in node.args)
            ):
                raise ValueError("another linear consumes the same input")
        blocks = [innermost_block(funcdef.body, call) for call in calls]
        if any(found is None for found in blocks):
            raise ValueError("projection calls not found in the function body")
        if len({id(block) for block, _ in blocks}) != 1:
            raise ValueError("projection calls are in different blocks")

        # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
        names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
        temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)]
        if names & set(temps):
            raise ValueError("fused temporaries would shadow existing names")
        merged = f"self.{self.merged_name}"
        sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
        template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)"
        assign = ast.parse(template).body[0]
        arg = next(
            node
            for node in ast.walk(assign)
            if isinstance(node, ast.Name) and node.id == "__arg__"
        )
        replace_expr(assign, arg, calls[0].args[0])
        block, index = blocks[0]
        ast.copy_location(assign, block[index])
        block.insert(min(index for _, index in blocks), assign)
        for call, temp in zip(calls, temps):
            replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))
        self.fused_forward = compile_forward(funcdef, fn)

    def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
        """Shapes must be compatible for a single merged, head-sharded GEMM."""
        q = module.get_submodule(self.q_name)
        k = module.get_submodule(self.k_name)
        v = module.get_submodule(self.v_name)
        head_size = model_config.get_head_size()
        compatible = (
            q.in_features == k.in_features == v.in_features
            and len({proj.bias is None for proj in (q, k, v)}) == 1
            and k.out_features == v.out_features
            and q.out_features % head_size == 0
            and k.out_features % head_size == 0
        )
        if not compatible:
            logger.debug("%s is not compatible with QKV fusion", type(module))
        return compatible

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