vllm.model_executor.models.transformers.fusers.moe ¶

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

class SharedExpertMLP(nn.Module):
    """Wraps an HF shared expert, applying the output gating it is paired with."""

    def __init__(self, shared_experts: nn.Module, gate: nn.Module | None = None):
        super().__init__()
        self.shared_experts = shared_experts
        self.gate = gate

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        out = self.shared_experts(hidden_states)
        if self.gate is not None:
            out = torch.sigmoid(self.gate(hidden_states)[0]) * out
        return out