forked from Karylab-cklius/vllm
99de48e98f
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
377 lines
13 KiB
Python
377 lines
13 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 ScaledRouter(TopKRouter):
|
|
"""Greedy router scaling its top-k weights (DeepSeek `routed_scaling_factor`)."""
|
|
|
|
def forward(self, hidden_states):
|
|
logits = F.linear(hidden_states, self.weight)
|
|
scores = F.softmax(logits, dim=-1)
|
|
value, index = torch.topk(scores, self.top_k, dim=-1)
|
|
value = value * 16.0
|
|
return logits, value, index
|
|
|
|
|
|
class Fp32Router(TopKRouter):
|
|
"""Router that computes its logits in fp32 (DeepSeek/GLM style)."""
|
|
|
|
def forward(self, hidden_states):
|
|
logits = F.linear(
|
|
hidden_states.type(torch.float32), self.weight.type(torch.float32)
|
|
)
|
|
scores = F.softmax(logits, dim=-1)
|
|
value, index = torch.topk(scores, self.top_k, dim=-1)
|
|
return logits, value.to(hidden_states.dtype), index
|
|
|
|
|
|
class GroupedRouter(TopKRouter):
|
|
"""Group-limited router (DeepSeek `group_limited_greedy`), scaled weights."""
|
|
|
|
def forward(self, hidden_states):
|
|
logits = F.linear(hidden_states, self.weight)
|
|
scores = F.softmax(logits, dim=-1)
|
|
group_scores = scores.view(-1, 4, 2).max(dim=-1).values
|
|
group_idx = torch.topk(group_scores, k=2, dim=-1)[1]
|
|
group_mask = torch.zeros_like(group_scores)
|
|
group_mask.scatter_(1, group_idx, 1)
|
|
score_mask = group_mask.unsqueeze(-1).expand(-1, 4, 2).reshape(-1, 8)
|
|
scores = scores.masked_fill(~score_mask.bool(), 0.0)
|
|
value, index = torch.topk(scores, self.top_k, dim=-1)
|
|
value = value * 16.0
|
|
return logits, value, index
|
|
|
|
|
|
class CorrectionRouter(nn.Module):
|
|
"""Router with a score-correction bias buffer (DeepSeek-V3 noaux) -> matched."""
|
|
|
|
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_matches_scaled_router():
|
|
"""Weight scaling after the top-k (DeepSeek style) does not break matching."""
|
|
with torch.device("meta"):
|
|
block = MoEBlock(ScaledRouter)
|
|
assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser)
|
|
|
|
|
|
def test_moe_fuser_matches_grouped_router():
|
|
"""Group masking between the score and the routing top-k still matches: the
|
|
scorer is anchored on the routing (last) top-k, not the group one."""
|
|
with torch.device("meta"):
|
|
block = MoEBlock(GroupedRouter)
|
|
fuser = MoEBlockFuser.match(block, "experts")
|
|
assert isinstance(fuser, MoEBlockFuser)
|
|
assert fuser.scoring_func == "softmax"
|
|
|
|
|
|
def test_moe_fuser_matches_correction_router():
|
|
"""A score-correction bias buffer (DeepSeek-V3 noaux) is allowed; the
|
|
rebuilt gate carries it in fp32 for FusedMoE's biased routers."""
|
|
with torch.device("meta"):
|
|
block = MoEBlock(CorrectionRouter)
|
|
fuser = MoEBlockFuser.match(block, "experts")
|
|
assert isinstance(fuser, MoEBlockFuser)
|
|
assert fuser.scoring_func == "sigmoid"
|
|
|
|
|
|
def test_moe_fuser_reads_router_dtype_from_the_gate():
|
|
"""A router that computes its logits in fp32 must keep routing in fp32 when
|
|
rebuilt, even though no config field names the dtype. The cast back to the
|
|
activation dtype after the top-k must not be mistaken for the routing dtype."""
|
|
with torch.device("meta"):
|
|
fp32 = MoEBlockFuser.match(MoEBlock(Fp32Router), "experts")
|
|
default = MoEBlockFuser.match(MoEBlock(TopKRouter), "experts")
|
|
assert fp32.router_dtype == torch.float32
|
|
assert default.router_dtype 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(BiasedRouter), # router with an unrecognized 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
|