| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """MoE fuser: route an HF MoE block through `FusedMoE` with vLLM's own routing.""" |
| 4 | |
| 5 | import ast |
| 6 | import inspect |
| 7 | import textwrap |
| 8 | import types |
| 9 | from collections.abc import Iterable, Iterator |
| 10 | from dataclasses import dataclass |
| 11 | |
| 12 | import torch |
| 13 | from torch import fx, nn |
| 14 | |
| 15 | from vllm.distributed import tensor_model_parallel_all_gather |
| 16 | from vllm.model_executor.layers.fused_moe import GateLinear |
| 17 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 18 | from vllm.model_executor.models.transformers.fx_utils import ( |
| 19 | find_node, |
| 20 | is_op, |
| 21 | peel, |
| 22 | trace, |
| 23 | ) |
| 24 | from vllm.model_executor.models.transformers.utils import named_state |
| 25 | from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk |
| 26 | |
| 27 | |
| 28 | def _own_returns(node: ast.AST) -> Iterator[ast.Return]: |
| 29 | """`return` statements in `node`'s own scope, not in nested functions.""" |
| 30 | stack = list(ast.iter_child_nodes(node)) |
| 31 | while stack: |
| 32 | child = stack.pop() |
| 33 | if isinstance(child, ast.Return): |
| 34 | yield child |
| 35 | elif not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): |
| 36 | stack.extend(ast.iter_child_nodes(child)) |
| 37 | |
| 38 | |
| 39 | def _returns_tuple(cls: type[nn.Module]) -> bool: |
| 40 | """Does `cls.forward()` return a tuple?""" |
| 41 | try: |
| 42 | source = textwrap.dedent(inspect.getsource(inspect.unwrap(cls.forward))) |
| 43 | forward = ast.parse(source).body[0] |
| 44 | except (OSError, SyntaxError, TypeError, IndexError): |
| 45 | return True |
| 46 | # Names bound to a tuple literal, e.g. `out = hidden, logits` then `return out`. |
| 47 | tuple_names = { |
| 48 | target.id |
| 49 | for node in ast.walk(forward) |
| 50 | if isinstance(node, ast.Assign) and isinstance(node.value, ast.Tuple) |
| 51 | for target in node.targets |
| 52 | if isinstance(target, ast.Name) |
| 53 | } |
| 54 | |
| 55 | def yields_tuple(value: ast.expr | None) -> bool: |
| 56 | if isinstance(value, ast.Tuple): |
| 57 | return True |
| 58 | if isinstance(value, ast.Name): |
| 59 | return value.id in tuple_names |
| 60 | if isinstance(value, ast.IfExp): |
| 61 | return yields_tuple(value.body) or yields_tuple(value.orelse) |
| 62 | return False |
| 63 | |
| 64 | return any(yields_tuple(ret.value) for ret in _own_returns(forward)) |
| 65 | |
| 66 | |
| 67 | def _is_scalar_gate(module: nn.Module) -> bool: |
| 68 | """A linear projecting to a single logit (the shared-expert sigmoid gate).""" |
| 69 | weight = getattr(module, "weight", None) |
| 70 | return ( |
| 71 | isinstance(module, nn.Linear) |
| 72 | and weight is not None |
| 73 | and weight.ndim == 2 |
| 74 | and weight.shape[0] == 1 |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | def _forced_dtype(nodes: Iterable[fx.Node]) -> torch.dtype | None: |
| 79 | """The floating dtype `nodes` cast to, if any. |
| 80 | |
| 81 | Computations that must run in higher precision say so in their forward (e.g. |
| 82 | `hidden_states.type(torch.float32)`), so the dtype is readable from the graph |
| 83 | even when the config does not name it.""" |
| 84 | for node in nodes: |
| 85 | if node.op not in ("call_method", "call_function"): |
| 86 | continue |
| 87 | name = str(node.target).rsplit(".", 1)[-1] |
| 88 | if name == "float": |
| 89 | return torch.float32 |
| 90 | if name in ("to", "type"): |
| 91 | for arg in (*node.args[1:], *node.kwargs.values()): |
| 92 | if isinstance(arg, torch.dtype) and arg.is_floating_point: |
| 93 | return arg |
| 94 | return None |
| 95 | |
| 96 | |
| 97 | def _reaches(node: fx.Node, key: str) -> set[fx.Node]: |
| 98 | """Returns the set of nodes reachable from `node` by following `key` edges.""" |
| 99 | seen: set[fx.Node] = set() |
| 100 | stack = [node] |
| 101 | while stack: |
| 102 | n = stack.pop() |
| 103 | if n in seen: |
| 104 | continue |
| 105 | seen.add(n) |
| 106 | stack.extend(getattr(n, key)) |
| 107 | return seen |
| 108 | |
| 109 | |
| 110 | class SharedExpertMLP(nn.Module): |
| 111 | """Wraps an HF shared expert, applying the output gating it is paired with.""" |
| 112 | |
| 113 | def __init__(self, shared_experts: nn.Module, gate: nn.Module | None = None): |
| 114 | super().__init__() |
| 115 | self.shared_experts = shared_experts |
| 116 | self.gate = gate |
| 117 | |
| 118 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 119 | out = self.shared_experts(hidden_states) |
| 120 | if self.gate is not None: |
| 121 | out = torch.sigmoid(self.gate(hidden_states)[0]) * out |
| 122 | return out |
| 123 | |
| 124 | |
| 125 | def _moe_block_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: |
| 126 | """Standard MoE block forward. |
| 127 | |
| 128 | Routing and any shared experts are handled inside `self.experts: MoERunner`.""" |
| 129 | orig_shape = hidden_states.shape |
| 130 | hidden_states = hidden_states.reshape(-1, orig_shape[-1]) |
| 131 | num_tokens = hidden_states.shape[0] |
| 132 | is_sequence_parallel = self.experts.moe_config.is_sequence_parallel |
| 133 | if is_sequence_parallel: |
| 134 | hidden_states = sequence_parallel_chunk(hidden_states) |
| 135 | out = self.experts(hidden_states, router_logits=hidden_states) |
| 136 | if is_sequence_parallel: |
| 137 | out = tensor_model_parallel_all_gather(out, 0)[:num_tokens] |
| 138 | return out.reshape(orig_shape) |
| 139 | |
| 140 | |
| 141 | @dataclass |
| 142 | class MoEBlockFuser: |
| 143 | """Fuser for MoE block `experts`, `gate` and `shared_experts` (optional).""" |
| 144 | |
| 145 | gate_name: str |
| 146 | scoring_func: str |
| 147 | shared_name: str | None |
| 148 | shared_gate_name: str | None |
| 149 | router_dtype: torch.dtype | None = None |
| 150 | |
| 151 | @staticmethod |
| 152 | def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: |
| 153 | """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`. |
| 154 | |
| 155 | Returns the scoring function and the dtype the router computes in.""" |
| 156 | state = {name for name, _ in named_state(gate)} |
| 157 | if "weight" not in state or state - {"weight", "e_score_correction_bias"}: |
| 158 | return None |
| 159 | graph = trace(gate) |
| 160 | if graph is None: |
| 161 | return None |
| 162 | # The routing top-k is the last one; any earlier one scores expert groups. |
| 163 | topks = [node for node in graph.nodes if is_op(node, "topk")] |
| 164 | if not topks: |
| 165 | return None |
| 166 | topk = topks[-1] |
| 167 | # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear. |
| 168 | scorers = [ |
| 169 | n |
| 170 | for n in _reaches(topk, "all_input_nodes") |
| 171 | if is_op(n, "softmax") or is_op(n, "sigmoid") |
| 172 | ] |
| 173 | if len(scorers) != 1: |
| 174 | return None |
| 175 | scorer = scorers[0] |
| 176 | logits_cone = _reaches(scorer, "all_input_nodes") |
| 177 | if not any(is_op(n, "linear") for n in logits_cone): |
| 178 | return None |
| 179 | scoring_func = "softmax" if is_op(scorer, "softmax") else "sigmoid" |
| 180 | return scoring_func, _forced_dtype(logits_cone) |
| 181 | |
| 182 | @staticmethod |
| 183 | def _match_shared_experts( |
| 184 | graph: fx.Graph, experts: str |
| 185 | ) -> tuple[str | None, str | None]: |
| 186 | """Detects the shared expert and its optional gate by dataflow.""" |
| 187 | experts_predicate = lambda n: n.op == "call_module" and n.target == experts |
| 188 | if (experts_node := find_node(graph, experts_predicate)) is None: |
| 189 | return None, None |
| 190 | from_experts = _reaches(experts_node, "users") |
| 191 | for add in graph.nodes: |
| 192 | if not is_op(add, "add"): |
| 193 | continue |
| 194 | operands = [a for a in add.args if isinstance(a, fx.Node)] |
| 195 | # Exactly one side is the experts' output; the other is the shared path. |
| 196 | sides = [a in from_experts for a in operands] |
| 197 | if len(operands) != 2 or sides.count(True) != 1: |
| 198 | continue |
| 199 | cone = _reaches(operands[sides.index(False)], "all_input_nodes") |
| 200 | modules = [n for n in cone if n.op == "call_module" and n.target != experts] |
| 201 | # A sigmoid wrapping one of those modules marks the shared-expert gate. |
| 202 | gate = next( |
| 203 | ( |
| 204 | src |
| 205 | for n in cone |
| 206 | if is_op(n, "sigmoid") |
| 207 | and isinstance(src := peel(n.args[0]), fx.Node) |
| 208 | and src in modules |
| 209 | ), |
| 210 | None, |
| 211 | ) |
| 212 | shared = [n for n in modules if n is not gate] |
| 213 | if len(shared) != 1: |
| 214 | return None, None |
| 215 | return shared[0].target, (gate.target if gate is not None else None) |
| 216 | return None, None |
| 217 | |
| 218 | @classmethod |
| 219 | def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None": |
| 220 | # Standard MoE block returns a single tensor. |
| 221 | if _returns_tuple(type(moe_block)): |
| 222 | return None |
| 223 | # Router: the child that scores + top-k selects. |
| 224 | gate_name = scoring_func = router_dtype = None |
| 225 | for name, child in moe_block.named_children(): |
| 226 | if ( |
| 227 | name != experts_name |
| 228 | and (router := cls._match_router(child)) is not None |
| 229 | ): |
| 230 | gate_name = name |
| 231 | scoring_func, router_dtype = router |
| 232 | break |
| 233 | if gate_name is None or scoring_func is None: |
| 234 | return None |
| 235 | # Shared expert: a child the block adds to the experts' output. |
| 236 | shared_name = shared_gate_name = None |
| 237 | others = [ |
| 238 | n |
| 239 | for n, _ in moe_block.named_children() |
| 240 | if n not in {experts_name, gate_name} |
| 241 | ] |
| 242 | if others: |
| 243 | graph = trace(moe_block) |
| 244 | if graph is None: |
| 245 | return None |
| 246 | shared_name, shared_gate_name = cls._match_shared_experts( |
| 247 | graph, experts_name |
| 248 | ) |
| 249 | if shared_gate_name is not None and not _is_scalar_gate( |
| 250 | getattr(moe_block, shared_gate_name) |
| 251 | ): |
| 252 | return None |
| 253 | # Fail closed: `rewrite_forward` runs only the experts and the detected |
| 254 | # shared expert, so any other stateful child would be dropped. |
| 255 | accounted = {experts_name, gate_name, shared_name, shared_gate_name} |
| 256 | for name, child in moe_block.named_children(): |
| 257 | if name not in accounted and next(named_state(child), None) is not None: |
| 258 | return None |
| 259 | return cls(gate_name, scoring_func, shared_name, shared_gate_name, router_dtype) |
| 260 | |
| 261 | def gate( |
| 262 | self, moe_block: nn.Module, prefix: str, out_dtype: torch.dtype | None = None |
| 263 | ) -> GateLinear: |
| 264 | """Rebuild the HF gate as a `GateLinear` for vLLM's fused MoE.""" |
| 265 | hf_gate = getattr(moe_block, self.gate_name) |
| 266 | num_experts, hidden_size = hf_gate.weight.shape |
| 267 | gate = GateLinear( |
| 268 | hidden_size, |
| 269 | num_experts, |
| 270 | bias=False, |
| 271 | out_dtype=out_dtype or self.router_dtype, |
| 272 | prefix=maybe_prefix(prefix, self.gate_name), |
| 273 | ) |
| 274 | if (bias := getattr(hf_gate, "e_score_correction_bias", None)) is not None: |
| 275 | gate.register_buffer("e_score_correction_bias", bias) |
| 276 | setattr(moe_block, self.gate_name, gate) |
| 277 | return gate |
| 278 | |
| 279 | def shared_experts( |
| 280 | self, moe_block: nn.Module, prefix: str |
| 281 | ) -> SharedExpertMLP | None: |
| 282 | """Build the HF shared expert (and its optional gate) |
| 283 | as a `SharedExpertMLP` for vLLM's fused MoE.""" |
| 284 | if self.shared_name is None: |
| 285 | return None |
| 286 | shared_experts = getattr(moe_block, self.shared_name) |
| 287 | gate = None |
| 288 | if self.shared_gate_name is not None: |
| 289 | hf_gate = getattr(moe_block, self.shared_gate_name) |
| 290 | gate = ReplicatedLinear( |
| 291 | hf_gate.in_features, |
| 292 | hf_gate.out_features, |
| 293 | bias=hf_gate.bias is not None, |
| 294 | prefix=maybe_prefix(prefix, self.shared_gate_name), |
| 295 | ) |
| 296 | setattr(moe_block, self.shared_gate_name, gate) |
| 297 | return SharedExpertMLP(shared_experts, gate) |
| 298 | |
| 299 | def rewrite_forward(self, moe_block: nn.Module) -> None: |
| 300 | """Rewrite `moe_block.forward` to route through vLLM's fused MoE.""" |
| 301 | moe_block.forward = types.MethodType(_moe_block_forward, moe_block) |