mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-09 06:25:01 +08:00
8a3cd90af5
Signed-off-by: Xin Yang <xyangx@amazon.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
77 lines
3.3 KiB
Python
77 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for the MoE grouped topk kernel
|
|
|
|
Run `pytest tests/kernels/moe/test_grouped_topk.py`.
|
|
"""
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.model_executor.layers.fused_moe.fused_moe import (fused_grouped_topk,
|
|
grouped_topk)
|
|
from vllm.platforms import current_platform
|
|
|
|
|
|
@pytest.mark.skipif(not current_platform.is_cuda(),
|
|
reason="This test is skipped on non-CUDA platform.")
|
|
@pytest.mark.parametrize("n_token", [1, 33, 64])
|
|
@pytest.mark.parametrize("n_hidden", [1024, 2048])
|
|
@pytest.mark.parametrize("n_expert", [16])
|
|
@pytest.mark.parametrize("topk", [2])
|
|
@pytest.mark.parametrize("renormalize", [True, False])
|
|
@pytest.mark.parametrize("num_expert_group", [8])
|
|
@pytest.mark.parametrize("topk_group", [2])
|
|
@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
|
|
@pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5])
|
|
@pytest.mark.parametrize("dtype",
|
|
[torch.float16, torch.bfloat16, torch.float32])
|
|
def test_grouped_topk(monkeypatch: pytest.MonkeyPatch, n_token: int,
|
|
n_hidden: int, n_expert: int, topk: int,
|
|
renormalize: bool, num_expert_group: int,
|
|
topk_group: int, scoring_func: str,
|
|
routed_scaling_factor: float, dtype: torch.dtype):
|
|
current_platform.seed_everything(0)
|
|
hidden_states = torch.randn((n_token, n_hidden),
|
|
dtype=dtype,
|
|
device="cuda")
|
|
gating_output = torch.randn((n_token, n_expert),
|
|
dtype=dtype,
|
|
device="cuda")
|
|
e_score_correction_bias = torch.randn((n_expert, ),
|
|
dtype=torch.float32,
|
|
device="cuda")
|
|
|
|
with monkeypatch.context() as m:
|
|
m.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0")
|
|
baseline_topk_weights, baseline_topk_ids = grouped_topk(
|
|
hidden_states=hidden_states,
|
|
gating_output=gating_output,
|
|
topk=topk,
|
|
renormalize=renormalize,
|
|
num_expert_group=num_expert_group,
|
|
topk_group=topk_group,
|
|
scoring_func=scoring_func,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
e_score_correction_bias=e_score_correction_bias)
|
|
|
|
test_topk_weights, test_topk_ids = fused_grouped_topk(
|
|
hidden_states=hidden_states,
|
|
gating_output=gating_output,
|
|
topk=topk,
|
|
renormalize=renormalize,
|
|
num_expert_group=num_expert_group,
|
|
topk_group=topk_group,
|
|
scoring_func=scoring_func,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
e_score_correction_bias=e_score_correction_bias)
|
|
|
|
if renormalize:
|
|
torch.testing.assert_close(baseline_topk_weights,
|
|
test_topk_weights,
|
|
atol=2e-2,
|
|
rtol=0)
|
|
torch.testing.assert_close(baseline_topk_ids,
|
|
test_topk_ids,
|
|
atol=0,
|
|
rtol=0)
|