144 lines
4.6 KiB
Python
144 lines
4.6 KiB
Python
import sys
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
|
|
from sglang.srt.layers.moe import hash_topk as hash_topk_module
|
|
from sglang.srt.layers.moe.hash_topk import HashTopK
|
|
from sglang.srt.layers.moe.topk import (
|
|
StandardTopKOutput,
|
|
)
|
|
from sglang.srt.models.deepseek_v2 import DeepseekV2MoE
|
|
from sglang.srt.runtime_context import get_parallel
|
|
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=5, suite="base-b-test-cpu")
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _set_dummy_server_args():
|
|
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
|
|
|
|
|
def test_hash_topk_remaps_per_rank_fused_shared_slots(monkeypatch):
|
|
monkeypatch.setattr(
|
|
hash_topk_module, "has_per_rank_fused_shared_slots", lambda *_args: True
|
|
)
|
|
recorded = {}
|
|
|
|
class FakeRecorder:
|
|
def on_select_experts(self, *, topk_ids):
|
|
recorded["topk_ids"] = topk_ids.clone()
|
|
|
|
from sglang.srt.runtime_context import get_resources
|
|
|
|
monkeypatch.setattr(get_resources(), "expert_distribution_recorder", FakeRecorder())
|
|
|
|
topk = HashTopK(
|
|
topk=3,
|
|
num_experts=256,
|
|
num_fused_shared_experts=1,
|
|
vocab_size=2,
|
|
scoring_func="sqrtsoftplus",
|
|
routed_scaling_factor=2.5,
|
|
)
|
|
with torch.no_grad():
|
|
topk.tid2eid.copy_(torch.tensor([[0, 65], [63, 127]], dtype=torch.int32))
|
|
|
|
info = ExpertLocationDispatchInfo(
|
|
ep_dispatch_algorithm="static",
|
|
partial_logical_to_rank_dispatch_physical_map=torch.arange(
|
|
256, dtype=torch.int32
|
|
),
|
|
partial_logical_to_all_physical_map=torch.arange(256, dtype=torch.int32).view(
|
|
256, 1
|
|
),
|
|
partial_logical_to_all_physical_map_num_valid=torch.ones(
|
|
256, dtype=torch.int32
|
|
),
|
|
num_physical_experts=256,
|
|
)
|
|
|
|
with (
|
|
get_parallel().override(moe_ep_size=4, moe_ep_rank=2),
|
|
hash_topk_module.envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.override(False),
|
|
):
|
|
output = topk(
|
|
hidden_states=torch.empty(2, 4),
|
|
router_logits=torch.ones(2, 256),
|
|
input_ids=torch.tensor([0, 1], dtype=torch.int64),
|
|
expert_location_dispatch_info=info,
|
|
)
|
|
|
|
# Physical layout for EP=4 has 64 routed slots per rank plus one local
|
|
# shared slot: [0..63, shared, 64..127, shared, ...].
|
|
assert output.topk_ids.tolist() == [[0, 66, 194], [63, 128, 194]]
|
|
assert torch.allclose(output.topk_weights[:, -1], torch.full((2,), 0.4))
|
|
assert recorded["topk_ids"].tolist() == [[0, 65], [63, 127]]
|
|
|
|
|
|
def test_hash_topk_empty_output_keeps_per_rank_shared_slot(monkeypatch):
|
|
monkeypatch.setattr(
|
|
hash_topk_module, "has_per_rank_fused_shared_slots", lambda *_args: True
|
|
)
|
|
|
|
topk = HashTopK(
|
|
topk=7,
|
|
num_experts=256,
|
|
num_fused_shared_experts=1,
|
|
vocab_size=2,
|
|
scoring_func="softmax",
|
|
)
|
|
|
|
output = topk.empty_topk_output(torch.device("cpu"))
|
|
|
|
assert output.topk_ids.shape == (0, 7)
|
|
assert output.topk_weights.shape == (0, 7)
|
|
assert output.router_logits.shape == (0, 6)
|
|
|
|
|
|
def test_deepep_empty_forward_does_not_append_shared_slot_twice():
|
|
captured = {}
|
|
|
|
class FakeTopK:
|
|
def empty_topk_output(self, device, *, layer_id=None):
|
|
return StandardTopKOutput(
|
|
topk_weights=torch.empty((0, 9), dtype=torch.float32, device=device),
|
|
topk_ids=torch.empty((0, 9), dtype=torch.int32, device=device),
|
|
router_logits=torch.empty((0, 8), dtype=torch.float32, device=device),
|
|
)
|
|
|
|
class FakeExperts:
|
|
should_fuse_routed_scaling_factor_in_topk = True
|
|
|
|
def __call__(self, hidden_states, topk_output):
|
|
captured["topk_ids_shape"] = tuple(topk_output.topk_ids.shape)
|
|
captured["topk_weights_shape"] = tuple(topk_output.topk_weights.shape)
|
|
return hidden_states
|
|
|
|
moe = SimpleNamespace(
|
|
_fuse_shared_experts_inside_sbo=False,
|
|
is_nextn=False,
|
|
num_fused_shared_experts=1,
|
|
layer_id=0,
|
|
topk=FakeTopK(),
|
|
experts=FakeExperts(),
|
|
alt_stream=None,
|
|
routed_scaling_factor=1.0,
|
|
)
|
|
|
|
hidden_states = torch.empty((0, 4), dtype=torch.float32)
|
|
forward_batch = SimpleNamespace(num_token_non_padded=None)
|
|
|
|
DeepseekV2MoE.forward_deepep(moe, hidden_states, forward_batch)
|
|
|
|
assert captured["topk_ids_shape"] == (0, 9)
|
|
assert captured["topk_weights_shape"] == (0, 9)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__, "-v"]))
|