wehub-resource-sync
300 lines
10 KiB
Python
300 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Unit tests for the Transformers modeling backend's MoE fuser."""
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from vllm.model_executor.models.transformers.fusers import MoEBlockFuser
|
|
|
|
from .test_linear import GLUMLP
|
|
|
|
|
|
class TopKRouter(nn.Module):
|
|
"""HF v5 top-k router: `linear -> softmax -> topk (-> renorm)`."""
|
|
|
|
def __init__(self, num_experts=8, hidden=16, top_k=2, sigmoid=False):
|
|
super().__init__()
|
|
self.top_k = top_k
|
|
self.sigmoid = sigmoid
|
|
self.weight = nn.Parameter(torch.zeros(num_experts, hidden))
|
|
|
|
def forward(self, hidden_states):
|
|
logits = F.linear(hidden_states, self.weight)
|
|
scores = torch.sigmoid(logits) if self.sigmoid else F.softmax(logits, dim=-1)
|
|
value, index = torch.topk(scores, self.top_k, dim=-1)
|
|
value = value / value.sum(dim=-1, keepdim=True)
|
|
return logits, value, index
|
|
|
|
|
|
class CorrectionRouter(nn.Module):
|
|
"""Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined."""
|
|
|
|
def __init__(self, num_experts=8, hidden=16):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.zeros(num_experts, hidden))
|
|
self.register_buffer("e_score_correction_bias", torch.zeros(num_experts))
|
|
|
|
def forward(self, hidden_states):
|
|
logits = F.linear(hidden_states, self.weight)
|
|
scores = torch.sigmoid(logits) + self.e_score_correction_bias
|
|
_, index = torch.topk(scores, 2, dim=-1)
|
|
return logits, scores, index
|
|
|
|
|
|
class BiasedRouter(TopKRouter):
|
|
"""A valid top-k router but not `weight`-only (extra `bias` param) -> declined."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.bias = nn.Parameter(torch.zeros(8))
|
|
|
|
def forward(self, hidden_states):
|
|
logits = F.linear(hidden_states, self.weight) + self.bias
|
|
scores = F.softmax(logits, dim=-1)
|
|
value, index = torch.topk(scores, self.top_k, dim=-1)
|
|
return logits, value, index
|
|
|
|
|
|
class DisconnectedRouter(TopKRouter):
|
|
"""linear+softmax+top-k present but top-k ignores the logits -> not a router."""
|
|
|
|
def forward(self, hidden_states):
|
|
logits = F.linear(hidden_states, self.weight)
|
|
_ = F.softmax(logits, dim=-1) # scored, but not consumed by top-k
|
|
value, index = torch.topk(hidden_states, self.top_k, dim=-1)
|
|
return logits, value, index
|
|
|
|
|
|
class MoEExperts(nn.Module):
|
|
"""Packed experts (3D weights); only its name (`experts`) matters here."""
|
|
|
|
def __init__(self, num_experts=8, hidden=16, inter=32):
|
|
super().__init__()
|
|
self.gate_up_proj = nn.Parameter(torch.zeros(num_experts, 2 * inter, hidden))
|
|
self.down_proj = nn.Parameter(torch.zeros(num_experts, hidden, inter))
|
|
|
|
def forward(self, hidden_states, index, weights):
|
|
return hidden_states
|
|
|
|
|
|
class MoEBlock(nn.Module):
|
|
"""Single-tensor MoE block (Qwen3-style); subclasses override `_shared`."""
|
|
|
|
def __init__(self, router_cls=TopKRouter):
|
|
super().__init__()
|
|
self.experts = MoEExperts()
|
|
self.gate = router_cls()
|
|
|
|
def _shared(self, x, logits):
|
|
"""The term added to the experts' output (none for a plain block)."""
|
|
return 0
|
|
|
|
def forward(self, hidden_states):
|
|
x = hidden_states.reshape(-1, hidden_states.shape[-1])
|
|
logits, weights, index = self.gate(x)
|
|
out = self.experts(x, index, weights) + self._shared(x, logits)
|
|
return out.reshape(hidden_states.shape)
|
|
|
|
|
|
class MoEBlockNoShared(MoEBlock):
|
|
"""No shared-expert child but a gate-derived add -> trace skipped, still fuses."""
|
|
|
|
def _shared(self, x, logits):
|
|
return logits.sum()
|
|
|
|
|
|
class MoEBlockShared(MoEBlock):
|
|
"""A block with a shared expert and its sigmoid gate (Qwen2-style)."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.shared_expert = GLUMLP()
|
|
self.shared_expert_gate = nn.Linear(16, 1, bias=False)
|
|
|
|
def _shared(self, x, logits):
|
|
return torch.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
|
|
|
|
|
|
class MoEBlockSharedNoGate(MoEBlock):
|
|
"""A block with an ungated shared expert -> native, shared passed through."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.shared_expert = GLUMLP()
|
|
|
|
def _shared(self, x, logits):
|
|
return self.shared_expert(x)
|
|
|
|
|
|
class MoEBlockTuple(MoEBlock):
|
|
"""A tuple-returning block (gpt-oss-style) -> must decline."""
|
|
|
|
def forward(self, hidden_states):
|
|
x = hidden_states.reshape(-1, hidden_states.shape[-1])
|
|
_, weights, index = self.gate(x)
|
|
return self.experts(x, index, weights), index
|
|
|
|
|
|
class MoEBlockTupleVar(MoEBlock):
|
|
"""Returns a name bound to a tuple, not a literal tuple -> must still decline."""
|
|
|
|
def forward(self, hidden_states):
|
|
x = hidden_states.reshape(-1, hidden_states.shape[-1])
|
|
_, weights, index = self.gate(x)
|
|
result = self.experts(x, index, weights), index
|
|
return result
|
|
|
|
|
|
class MoEBlockNestedTupleReturn(MoEBlock):
|
|
"""Tuple `return` in a nested helper; block returns one tensor -> still fuses."""
|
|
|
|
def forward(self, hidden_states):
|
|
def keep(a, b):
|
|
return a, b
|
|
|
|
x = hidden_states.reshape(-1, hidden_states.shape[-1])
|
|
_, weights, index = self.gate(x)
|
|
out, _ = keep(self.experts(x, index, weights), index)
|
|
return out.reshape(hidden_states.shape)
|
|
|
|
|
|
class PlainMLP(nn.Module):
|
|
"""A non-GLU FFN: `down(act(up(x)))`, no gating multiply."""
|
|
|
|
def __init__(self, hidden: int = 16, inter: int = 32):
|
|
super().__init__()
|
|
self.up_proj = nn.Linear(hidden, inter, bias=False)
|
|
self.down_proj = nn.Linear(inter, hidden, bias=False)
|
|
self.act_fn = nn.SiLU()
|
|
|
|
def forward(self, x):
|
|
return self.down_proj(self.act_fn(self.up_proj(x)))
|
|
|
|
|
|
class MoEBlockSharedNonGLU(MoEBlock):
|
|
"""A non-GLU shared expert -> detected by dataflow (no gate/up merge)."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.shared_expert = PlainMLP()
|
|
|
|
def _shared(self, x, logits):
|
|
return self.shared_expert(x)
|
|
|
|
|
|
class MoEBlockUnaccounted(MoEBlock):
|
|
"""A weight-bearing child outside the fused dataflow (pre-router) -> declined."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.extra = nn.Linear(16, 16, bias=False)
|
|
|
|
def forward(self, hidden_states):
|
|
x = self.extra(hidden_states.reshape(-1, hidden_states.shape[-1]))
|
|
_, weights, index = self.gate(x)
|
|
return self.experts(x, index, weights).reshape(hidden_states.shape)
|
|
|
|
|
|
class BufferScale(nn.Module):
|
|
"""A stateful child carrying only a buffer (no parameters)."""
|
|
|
|
def __init__(self, hidden: int = 16):
|
|
super().__init__()
|
|
self.register_buffer("scale", torch.ones(hidden))
|
|
|
|
def forward(self, x):
|
|
return x * self.scale
|
|
|
|
|
|
class MoEBlockUnaccountedBuffer(MoEBlockUnaccounted):
|
|
"""Like `MoEBlockUnaccounted`, but the extra child holds only a buffer."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.extra = BufferScale()
|
|
|
|
|
|
@pytest.mark.parametrize("sigmoid", [False, True])
|
|
def test_moe_fuser_detects_router(sigmoid):
|
|
with torch.device("meta"):
|
|
block = MoEBlock(lambda: TopKRouter(sigmoid=sigmoid))
|
|
fuser = MoEBlockFuser.match(block, "experts")
|
|
assert isinstance(fuser, MoEBlockFuser)
|
|
assert fuser.gate_name == "gate"
|
|
assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax")
|
|
assert fuser.shared_name is None and fuser.shared_gate_name is None
|
|
|
|
|
|
def test_moe_fuser_detects_shared_experts():
|
|
with torch.device("meta"):
|
|
block = MoEBlockShared()
|
|
fuser = MoEBlockFuser.match(block, "experts")
|
|
assert isinstance(fuser, MoEBlockFuser)
|
|
assert fuser.shared_name == "shared_expert"
|
|
assert fuser.shared_gate_name == "shared_expert_gate"
|
|
|
|
|
|
def test_moe_fuser_skips_shared_detection_without_extra_children():
|
|
"""With only experts and gate, shared-expert detection (and its block trace)
|
|
is skipped, so a gate-derived add is not misread as a shared expert."""
|
|
with torch.device("meta"):
|
|
block = MoEBlockNoShared()
|
|
fuser = MoEBlockFuser.match(block, "experts")
|
|
assert isinstance(fuser, MoEBlockFuser)
|
|
assert fuser.shared_name is None and fuser.shared_gate_name is None
|
|
|
|
|
|
def test_moe_fuser_shared_without_gate():
|
|
with torch.device("meta"):
|
|
block = MoEBlockSharedNoGate()
|
|
fuser = MoEBlockFuser.match(block, "experts")
|
|
assert isinstance(fuser, MoEBlockFuser)
|
|
assert fuser.shared_name == "shared_expert"
|
|
assert fuser.shared_gate_name is None
|
|
|
|
|
|
def test_moe_fuser_detects_non_glu_shared_expert():
|
|
with torch.device("meta"):
|
|
block = MoEBlockSharedNonGLU()
|
|
fuser = MoEBlockFuser.match(block, "experts")
|
|
assert isinstance(fuser, MoEBlockFuser)
|
|
# Recognised by dataflow (added to the experts' output), though not a GLU.
|
|
assert fuser.shared_name == "shared_expert"
|
|
assert fuser.shared_gate_name is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"block_cls",
|
|
[
|
|
lambda: MoEBlock(CorrectionRouter), # score-correction buffer (grouped)
|
|
lambda: MoEBlock(BiasedRouter), # router not weight-only (extra param)
|
|
MoEBlockTuple, # tuple-returning block (e.g. gpt-oss)
|
|
MoEBlockTupleVar, # tuple returned via a name binding, not a literal
|
|
MoEBlockUnaccounted, # weight-bearing child outside the fused dataflow
|
|
MoEBlockUnaccountedBuffer, # buffer-only child outside the fused dataflow
|
|
],
|
|
)
|
|
def test_moe_fuser_declines_unsupported(block_cls):
|
|
with torch.device("meta"):
|
|
block = block_cls()
|
|
assert MoEBlockFuser.match(block, "experts") is None
|
|
|
|
|
|
def test_moe_fuser_ignores_nested_returns():
|
|
"""A tuple `return` inside a nested helper must not decline a block whose own
|
|
forward returns a single tensor."""
|
|
with torch.device("meta"):
|
|
block = MoEBlockNestedTupleReturn()
|
|
assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser)
|
|
|
|
|
|
def test_moe_fuser_router_requires_connected_dataflow():
|
|
"""A gate with linear + softmax + top-k present but not wired as a router
|
|
(top-k selects over the input, not the scored logits) is not detected."""
|
|
with torch.device("meta"):
|
|
block = MoEBlock(DisconnectedRouter)
|
|
assert MoEBlockFuser.match(block, "experts") is None
|