711 lines
28 KiB
Python
711 lines
28 KiB
Python
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torch_npu
|
|
from pytest_mock import MockerFixture
|
|
|
|
from tests.ut.base import TestBase
|
|
from vllm_ascend.ascend_forward_context import MoECommType
|
|
from vllm_ascend.ops.fused_moe.experts_selector import select_experts, zero_experts_compute
|
|
from vllm_ascend.ops.fused_moe.moe_runtime_args import MoEPrepareOutput
|
|
from vllm_ascend.utils import AscendDeviceType, adapt_patch, enable_custom_op
|
|
|
|
adapt_patch(True)
|
|
|
|
|
|
def mock_ep_and_mc2_group(mocker):
|
|
mock_group = mocker.MagicMock()
|
|
mock_group.rank_in_group = 0
|
|
mock_group.rank = 0
|
|
mock_group.world_size = 4
|
|
mock_group.device_group = "mock_group_ep"
|
|
mock_group.all_to_all = MagicMock(return_value=torch.randn(8, 8))
|
|
return mock_group
|
|
|
|
|
|
def mock_dp_and_tp_group(mocker):
|
|
mock_group = mocker.MagicMock()
|
|
mock_group.rank_in_group = 0
|
|
mock_group.world_size = 2
|
|
mock_group.device_group = "mock_group"
|
|
mock_group.all_gather = MagicMock(return_value=torch.randn(10, 32))
|
|
return mock_group
|
|
|
|
|
|
def _require_ascend_custom_op(op_name: str):
|
|
try:
|
|
custom_op_enabled = enable_custom_op()
|
|
except Exception as exc:
|
|
pytest.skip(f"requires vllm_ascend custom ops: {exc}")
|
|
if not custom_op_enabled:
|
|
pytest.skip("requires vllm_ascend custom ops")
|
|
try:
|
|
getattr(torch.ops._C_ascend, op_name)
|
|
except AttributeError:
|
|
pytest.skip(f"requires torch.ops._C_ascend.{op_name}")
|
|
|
|
|
|
def _sort_topk_by_ids(topk_weights: torch.Tensor, topk_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
sorted_ids, order = torch.sort(topk_ids.to(torch.int64), dim=-1)
|
|
sorted_weights = topk_weights.gather(1, order)
|
|
return sorted_weights, sorted_ids.to(torch.int32)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def setup_vllm_config_mock(mocker: MockerFixture):
|
|
mock_hf_config = MagicMock()
|
|
mock_hf_config.model_type = "llama"
|
|
|
|
mock_model_config = MagicMock()
|
|
mock_model_config.hf_config = mock_hf_config
|
|
|
|
mock_vllm_config = MagicMock()
|
|
mock_vllm_config.model_config = mock_model_config
|
|
mock_vllm_config.parallel_config = MagicMock(tensor_parallel_size=2)
|
|
mock_vllm_config.scheduler_config = MagicMock(max_num_seqs=4)
|
|
mock_vllm_config.model_config.max_model_len = 2048
|
|
|
|
mocker.patch("vllm_ascend.ops.fused_moe.fused_moe.get_current_vllm_config", return_value=mock_vllm_config)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_dist_env(mocker: MockerFixture):
|
|
mock_moe_comm_method = MagicMock()
|
|
|
|
def mock_prepare(hidden_states, router_logits, **kwargs):
|
|
return MoEPrepareOutput(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
mc2_mask=kwargs.get("mc2_mask"),
|
|
padded_hidden_states_shape=None,
|
|
pertoken_scale=None,
|
|
)
|
|
|
|
mock_moe_comm_method.prepare.side_effect = mock_prepare
|
|
|
|
mock_fused_experts_result = torch.randn(16, 2)
|
|
mock_moe_comm_method.fused_experts.return_value = mock_fused_experts_result
|
|
|
|
def mock_finalize(hidden_states, **kwargs):
|
|
return hidden_states
|
|
|
|
mock_moe_comm_method.finalize.side_effect = mock_finalize
|
|
dp_metadata = MagicMock(num_tokens_across_dp_cpu=[5, 5])
|
|
mock_weight_prefetch_method = MagicMock()
|
|
mock_forward_context_obj = MagicMock(
|
|
moe_comm_method=mock_moe_comm_method,
|
|
moe_comm_type=MoECommType.MC2,
|
|
max_tokens_across_dp=10,
|
|
dp_metadata=dp_metadata,
|
|
mc2_mask=torch.zeros(16, dtype=torch.bool),
|
|
padded_num_tokens=16,
|
|
with_quant=False,
|
|
)
|
|
|
|
with (
|
|
patch("torch.distributed.get_rank", return_value=0),
|
|
patch("torch.distributed.get_world_size", return_value=4),
|
|
patch("vllm_ascend.ops.fused_moe.fused_moe.get_ep_group", return_value=mock_ep_and_mc2_group(mocker)),
|
|
patch("vllm_ascend.ops.fused_moe.token_dispatcher.get_ep_group", return_value=mock_ep_and_mc2_group(mocker)),
|
|
patch("vllm_ascend.ops.fused_moe.fused_moe.get_mc2_group", return_value=mock_ep_and_mc2_group(mocker)),
|
|
patch("vllm_ascend.ops.fused_moe.fused_moe.get_tp_group", return_value=mock_dp_and_tp_group(mocker)),
|
|
patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_dp_and_tp_group(mocker)),
|
|
patch("vllm_ascend.ops.fused_moe.fused_moe.get_dp_group", return_value=mock_dp_and_tp_group(mocker)),
|
|
patch("vllm.model_executor.layers.fused_moe.layer.get_dp_group", return_value=mock_dp_and_tp_group(mocker)),
|
|
patch("vllm.model_executor.layers.fused_moe.config.get_dp_group", return_value=mock_dp_and_tp_group(mocker)),
|
|
patch(
|
|
"vllm_ascend.ops.fused_moe.fused_moe.get_ascend_config",
|
|
return_value=MagicMock(enable_multistream_moe=False, expert_map_path=None),
|
|
),
|
|
patch(
|
|
"vllm_ascend.ops.fused_moe.fused_moe.init_eplb_config",
|
|
return_value=(torch.tensor([0, 1, 2, -1, -1, -1, -1, -1]), None, 0),
|
|
),
|
|
patch("vllm_ascend.ops.fused_moe.fused_moe.get_forward_context", return_value=mock_forward_context_obj),
|
|
patch("vllm_ascend.ascend_forward_context.get_forward_context", return_value=mock_forward_context_obj),
|
|
patch("vllm_ascend.utils.get_ascend_device_type", return_value=AscendDeviceType.A3),
|
|
patch("vllm_ascend.ops.fused_moe.moe_comm_method.MC2CommImpl._get_token_dispatcher", return_value=None),
|
|
patch("vllm_ascend.ops.fused_moe.moe_comm_method.AlltoAllCommImpl._get_token_dispatcher", return_value=None),
|
|
patch("vllm_ascend.ops.fused_moe.moe_comm_method.AllGatherCommImpl._get_token_dispatcher", return_value=None),
|
|
patch(
|
|
"vllm_ascend.ops.fused_moe.experts_selector.get_weight_prefetch_method",
|
|
return_value=mock_weight_prefetch_method,
|
|
),
|
|
):
|
|
yield {
|
|
"mock_forward_context_obj": mock_forward_context_obj,
|
|
"mock_moe_comm_method": mock_moe_comm_method,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_moe_env(mocker: MockerFixture):
|
|
with (
|
|
patch(
|
|
"torch_npu.npu_moe_init_routing",
|
|
return_value=(torch.randn(8, 2), torch.randint(0, 8, (8, 2)), torch.tensor([0, 1, 2, 4, 6, 2, 7, 1])),
|
|
),
|
|
patch("torch_npu.npu_moe_compute_expert_tokens", return_value=(torch.randn(8, 2))),
|
|
patch("torch_npu.npu_moe_distribute_dispatch", return_value=(torch.randn(16, 2))),
|
|
patch("torch_npu.npu_moe_distribute_combine", return_value=(torch.randn(16, 2))),
|
|
patch("torch_npu.npu_grouped_matmul", return_value=([torch.randn(16, 2)])),
|
|
patch("torch_npu.npu_swiglu", return_value=(torch.randn(16, 2))),
|
|
patch("torch_npu.npu_moe_finalize_routing", return_value=(torch.randn(16, 2))),
|
|
):
|
|
if hasattr(torch_npu, "npu_moe_distribute_dispatch_v2"):
|
|
with (
|
|
patch("torch_npu.npu_moe_distribute_dispatch_v2", return_value=(torch.randn(16, 2))),
|
|
patch("torch_npu.npu_moe_distribute_combine_v2", return_value=(torch.randn(16, 2))),
|
|
):
|
|
yield
|
|
else:
|
|
yield
|
|
|
|
|
|
class TestExpertsSelector:
|
|
@pytest.mark.parametrize("num_experts", [256, 128])
|
|
def test_select_experts(self, mock_dist_env, mock_moe_env, num_experts):
|
|
x = torch.randn(8, 2)
|
|
router_logits = torch.randn(8, 2)
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=x,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=False,
|
|
renormalize=True,
|
|
topk_group=None,
|
|
num_expert_group=None,
|
|
custom_routing_function=None,
|
|
scoring_func="softmax",
|
|
e_score_correction_bias=None,
|
|
num_experts=num_experts,
|
|
)
|
|
|
|
assert topk_weights.shape == (8, 2)
|
|
assert topk_ids.shape == (8, 2)
|
|
|
|
@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
|
|
@pytest.mark.parametrize("renormalize", [True, False])
|
|
def test_select_experts_with_different_scoring_func(self, mock_dist_env, mock_moe_env, scoring_func, renormalize):
|
|
num_tokens = 16
|
|
num_experts = 8
|
|
hidden_size = 32
|
|
|
|
hidden_states = torch.randn(num_tokens, hidden_size)
|
|
router_logits = torch.randn(num_tokens, num_experts)
|
|
|
|
def simple_custom_routing(hidden_states, gating_output, topk, renormalize):
|
|
if scoring_func == "softmax":
|
|
weights = gating_output.softmax(dim=-1)
|
|
else:
|
|
weights = gating_output.sigmoid()
|
|
topk_weights, topk_ids = weights.topk(topk, dim=-1)
|
|
if renormalize:
|
|
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
|
return topk_weights, topk_ids.to(torch.int32)
|
|
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=False,
|
|
renormalize=renormalize,
|
|
topk_group=None,
|
|
num_expert_group=None,
|
|
custom_routing_function=simple_custom_routing,
|
|
scoring_func=scoring_func,
|
|
e_score_correction_bias=None,
|
|
num_experts=num_experts,
|
|
)
|
|
|
|
assert topk_weights.shape == (num_tokens, 2)
|
|
assert topk_ids.shape == (num_tokens, 2)
|
|
assert topk_weights.dtype == hidden_states.dtype
|
|
assert topk_ids.dtype == torch.int32
|
|
|
|
if renormalize:
|
|
weight_sum = topk_weights.sum(dim=-1)
|
|
torch.testing.assert_close(weight_sum, torch.ones_like(weight_sum), rtol=1e-4, atol=1e-4)
|
|
|
|
def test_select_experts_with_grouped_topk(self, mock_dist_env, mock_moe_env):
|
|
_require_ascend_custom_op("moe_gating_top_k")
|
|
|
|
num_tokens = 16
|
|
num_experts = 8
|
|
hidden_size = 32
|
|
num_expert_group = 4
|
|
topk_group = 2
|
|
|
|
hidden_states = torch.randn(num_tokens, hidden_size, device="npu")
|
|
router_logits = torch.randn(num_tokens, num_experts, device="npu")
|
|
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=True,
|
|
renormalize=True,
|
|
topk_group=topk_group,
|
|
num_expert_group=num_expert_group,
|
|
custom_routing_function=None,
|
|
scoring_func="softmax",
|
|
e_score_correction_bias=None,
|
|
num_experts=num_experts,
|
|
)
|
|
|
|
assert topk_weights.shape == (num_tokens, 2)
|
|
assert topk_ids.shape == (num_tokens, 2)
|
|
|
|
def test_select_experts_with_e_score_correction_bias(self, mock_dist_env, mock_moe_env):
|
|
_require_ascend_custom_op("moe_gating_top_k")
|
|
|
|
num_tokens = 16
|
|
num_experts = 8
|
|
hidden_size = 32
|
|
|
|
hidden_states = torch.randn(num_tokens, hidden_size, device="npu")
|
|
router_logits = torch.randn(num_tokens, num_experts, device="npu")
|
|
e_score_correction_bias = torch.randn(num_experts, device="npu")
|
|
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=False,
|
|
renormalize=True,
|
|
topk_group=None,
|
|
num_expert_group=None,
|
|
custom_routing_function=None,
|
|
scoring_func="softmax",
|
|
e_score_correction_bias=e_score_correction_bias,
|
|
num_experts=num_experts,
|
|
)
|
|
|
|
assert topk_weights.shape == (num_tokens, 2)
|
|
assert topk_ids.shape == (num_tokens, 2)
|
|
|
|
def test_select_experts_with_custom_routing_function(self, mock_dist_env, mock_moe_env):
|
|
num_tokens = 16
|
|
num_experts = 8
|
|
hidden_size = 32
|
|
|
|
hidden_states = torch.randn(num_tokens, hidden_size)
|
|
router_logits = torch.randn(num_tokens, num_experts)
|
|
|
|
def custom_routing(hidden_states, gating_output, topk, renormalize):
|
|
weights = gating_output.softmax(dim=-1)
|
|
topk_weights, topk_ids = weights.topk(topk, dim=-1)
|
|
if renormalize:
|
|
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
|
return topk_weights, topk_ids.to(torch.int32)
|
|
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=False,
|
|
renormalize=True,
|
|
topk_group=None,
|
|
num_expert_group=None,
|
|
custom_routing_function=custom_routing,
|
|
scoring_func="softmax",
|
|
e_score_correction_bias=None,
|
|
num_experts=num_experts,
|
|
)
|
|
|
|
assert topk_weights.shape == (num_tokens, 2)
|
|
assert topk_ids.shape == (num_tokens, 2)
|
|
|
|
def test_select_experts_weight_sum_range(self, mock_dist_env, mock_moe_env):
|
|
num_tokens = 16
|
|
num_experts = 8
|
|
hidden_size = 32
|
|
|
|
hidden_states = torch.randn(num_tokens, hidden_size)
|
|
router_logits = torch.randn(num_tokens, num_experts)
|
|
|
|
def simple_routing(hidden_states, gating_output, topk, renormalize):
|
|
weights = gating_output.softmax(dim=-1)
|
|
topk_weights, topk_ids = weights.topk(topk, dim=-1)
|
|
if renormalize:
|
|
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
|
return topk_weights, topk_ids.to(torch.int32)
|
|
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=False,
|
|
renormalize=True,
|
|
topk_group=None,
|
|
num_expert_group=None,
|
|
custom_routing_function=simple_routing,
|
|
scoring_func="softmax",
|
|
e_score_correction_bias=None,
|
|
num_experts=num_experts,
|
|
)
|
|
|
|
assert (topk_weights >= 0).all()
|
|
assert (topk_weights <= 1).all()
|
|
|
|
weight_sum = topk_weights.sum(dim=-1)
|
|
torch.testing.assert_close(weight_sum, torch.ones_like(weight_sum), rtol=1e-4, atol=1e-4)
|
|
|
|
def test_select_experts_expert_id_range(self, mock_dist_env, mock_moe_env):
|
|
num_tokens = 16
|
|
num_experts = 8
|
|
hidden_size = 32
|
|
|
|
hidden_states = torch.randn(num_tokens, hidden_size)
|
|
router_logits = torch.randn(num_tokens, num_experts)
|
|
|
|
def simple_routing(hidden_states, gating_output, topk, renormalize):
|
|
weights = gating_output.softmax(dim=-1)
|
|
topk_weights, topk_ids = weights.topk(topk, dim=-1)
|
|
if renormalize:
|
|
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
|
return topk_weights, topk_ids.to(torch.int32)
|
|
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=False,
|
|
renormalize=True,
|
|
topk_group=None,
|
|
num_expert_group=None,
|
|
custom_routing_function=simple_routing,
|
|
scoring_func="softmax",
|
|
e_score_correction_bias=None,
|
|
num_experts=num_experts,
|
|
)
|
|
|
|
assert (topk_ids >= 0).all()
|
|
assert (topk_ids < num_experts).all()
|
|
|
|
@patch("vllm_ascend.ops.fused_moe.experts_selector.get_weight_prefetch_method")
|
|
@patch("vllm_ascend.ops.fused_moe.experts_selector.check_npu_moe_gating_top_k", return_value=False)
|
|
def test_select_experts_native_softmax_matches_expected(self, _, mock_get_weight_prefetch_method):
|
|
hidden_states = torch.tensor([[1.0, 0.0, -1.0, 2.0], [0.5, 1.5, -0.5, 1.0]], dtype=torch.float32)
|
|
router_logits = torch.tensor([[3.0, 1.0, 0.0, 2.0], [0.0, 4.0, 1.0, 2.0]], dtype=torch.float32)
|
|
|
|
prefetch_method = MagicMock()
|
|
mock_get_weight_prefetch_method.return_value = prefetch_method
|
|
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=False,
|
|
renormalize=True,
|
|
topk_group=None,
|
|
num_expert_group=None,
|
|
custom_routing_function=None,
|
|
scoring_func="softmax",
|
|
e_score_correction_bias=None,
|
|
num_experts=4,
|
|
)
|
|
|
|
expected_probs = router_logits.softmax(dim=-1)
|
|
expected_weights, expected_ids = expected_probs.topk(2, dim=-1)
|
|
expected_weights = expected_weights / expected_weights.sum(dim=-1, keepdim=True)
|
|
|
|
prefetch_method.maybe_prefetch_moe_weight_preprocess.assert_called_once_with(hidden_states, "gate_up")
|
|
torch.testing.assert_close(topk_weights, expected_weights.to(hidden_states.dtype))
|
|
assert torch.equal(topk_ids, expected_ids.to(torch.int32))
|
|
|
|
def test_select_experts_grouped_topk_bias_uses_original_weights(self):
|
|
_require_ascend_custom_op("moe_gating_top_k")
|
|
|
|
hidden_states = torch.tensor([[1.0, 0.0, -1.0, 2.0]], dtype=torch.float32, device="npu")
|
|
router_logits = torch.tensor([[4.0, 3.0, 1.0, 0.0]], dtype=torch.float32, device="npu")
|
|
e_score_correction_bias = torch.tensor([-10.0, -10.0, 5.0, 5.0], dtype=torch.float32, device="npu")
|
|
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=True,
|
|
renormalize=False,
|
|
topk_group=1,
|
|
num_expert_group=2,
|
|
custom_routing_function=None,
|
|
scoring_func="softmax",
|
|
e_score_correction_bias=e_score_correction_bias,
|
|
num_experts=4,
|
|
)
|
|
|
|
original_weights = router_logits.softmax(dim=-1)
|
|
actual_weights, actual_ids = _sort_topk_by_ids(topk_weights.cpu(), topk_ids.cpu())
|
|
expected_ids = torch.tensor([[2, 3]], dtype=torch.int32)
|
|
expected_weights = original_weights[:, 2:4].to("cpu")
|
|
|
|
assert torch.equal(actual_ids, expected_ids)
|
|
torch.testing.assert_close(actual_weights, expected_weights)
|
|
|
|
@patch("vllm_ascend.ops.fused_moe.experts_selector.get_weight_prefetch_method", return_value=MagicMock())
|
|
@patch("vllm_ascend.ops.fused_moe.experts_selector.check_npu_moe_gating_top_k", return_value=False)
|
|
def test_select_experts_invalid_scoring_func_raises(self, _, __):
|
|
hidden_states = torch.randn(2, 4)
|
|
router_logits = torch.randn(2, 4)
|
|
|
|
with pytest.raises(ValueError, match="Unsupported scoring function"):
|
|
select_experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=False,
|
|
renormalize=True,
|
|
topk_group=None,
|
|
num_expert_group=None,
|
|
custom_routing_function=None,
|
|
scoring_func="unsupported",
|
|
e_score_correction_bias=None,
|
|
num_experts=4,
|
|
)
|
|
|
|
|
|
class TestFusedMoENPUOpsAccuracy:
|
|
def test_moe_gating_top_k_matches_native_softmax(self, mock_dist_env):
|
|
_require_ascend_custom_op("moe_gating_top_k")
|
|
|
|
hidden_states = torch.randn(4, 16, dtype=torch.float32, device="npu")
|
|
router_logits_cpu = torch.tensor(
|
|
[[3.0, 1.0, 0.0, 2.0], [0.0, 4.0, 1.0, 2.0], [1.0, 0.5, 5.0, -1.0], [2.5, -0.5, 1.5, 0.25]],
|
|
dtype=torch.float32,
|
|
)
|
|
router_logits = router_logits_cpu.to("npu")
|
|
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
top_k=2,
|
|
use_grouped_topk=False,
|
|
renormalize=True,
|
|
topk_group=None,
|
|
num_expert_group=None,
|
|
custom_routing_function=None,
|
|
scoring_func="softmax",
|
|
e_score_correction_bias=None,
|
|
num_experts=4,
|
|
)
|
|
|
|
expected_weights, expected_ids = router_logits_cpu.softmax(dim=-1).topk(2, dim=-1)
|
|
expected_weights = expected_weights / expected_weights.sum(dim=-1, keepdim=True)
|
|
actual_weights, actual_ids = _sort_topk_by_ids(topk_weights.cpu(), topk_ids.cpu())
|
|
expected_weights, expected_ids = _sort_topk_by_ids(expected_weights, expected_ids.to(torch.int32))
|
|
|
|
assert torch.equal(actual_ids, expected_ids)
|
|
torch.testing.assert_close(actual_weights, expected_weights, rtol=1e-3, atol=1e-3)
|
|
|
|
def test_npu_swiglu_matches_torch_reference(self):
|
|
gate_up_cpu = torch.tensor(
|
|
[[1.0, -2.0, 0.5, 3.0], [-0.5, 2.0, 1.5, -1.0], [4.0, -3.0, -2.0, 0.25]],
|
|
dtype=torch.float16,
|
|
)
|
|
gate_up = gate_up_cpu.to("npu")
|
|
|
|
actual = torch_npu.npu_swiglu(gate_up)
|
|
left, right = gate_up_cpu.chunk(2, dim=-1)
|
|
expected = F.silu(left.float()) * right.float()
|
|
|
|
torch.testing.assert_close(actual.cpu().float(), expected, rtol=2e-3, atol=2e-3)
|
|
|
|
def test_npu_grouped_matmul_matches_torch_reference(self):
|
|
hidden_states_cpu = torch.tensor(
|
|
[[1.0, 2.0, -1.0], [0.5, -0.5, 1.0], [2.0, 0.0, 1.0], [-1.0, 3.0, 0.5]],
|
|
dtype=torch.float16,
|
|
)
|
|
weight_cpu = torch.tensor(
|
|
[[[1.0, 0.0], [0.5, -1.0], [-0.5, 2.0]], [[-1.0, 1.5], [2.0, 0.25], [0.0, -0.5]]],
|
|
dtype=torch.float16,
|
|
)
|
|
group_list = torch.tensor([2, 2], dtype=torch.int64, device="npu")
|
|
|
|
actual = torch_npu.npu_grouped_matmul(
|
|
x=[hidden_states_cpu.to("npu")],
|
|
weight=[weight_cpu.to("npu")],
|
|
split_item=3,
|
|
group_list_type=1,
|
|
group_type=0,
|
|
group_list=group_list,
|
|
output_dtype=torch.float16,
|
|
)[0]
|
|
expected = torch.cat(
|
|
[
|
|
hidden_states_cpu[:2].float() @ weight_cpu[0].float(),
|
|
hidden_states_cpu[2:].float() @ weight_cpu[1].float(),
|
|
],
|
|
dim=0,
|
|
)
|
|
|
|
torch.testing.assert_close(actual.cpu().float(), expected, rtol=2e-3, atol=2e-3)
|
|
|
|
|
|
class TestZeroExpertsCompute(TestBase):
|
|
def test_zero_experts_compute_identity_type(self):
|
|
num_experts = 8
|
|
num_tokens = 4
|
|
top_k = 2
|
|
hidden_size = 16
|
|
|
|
expert_indices = torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7]], dtype=torch.int32)
|
|
expert_scales = torch.ones(num_tokens, top_k, dtype=torch.float32)
|
|
hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.float32)
|
|
|
|
result_indices, result_scales, result_hidden = zero_experts_compute(
|
|
expert_indices=expert_indices,
|
|
expert_scales=expert_scales,
|
|
num_experts=num_experts,
|
|
zero_expert_type="identity",
|
|
hidden_states=hidden_states,
|
|
)
|
|
|
|
assert result_indices.shape == expert_indices.shape
|
|
assert result_scales.shape == expert_scales.shape
|
|
assert result_hidden.shape == hidden_states.shape
|
|
|
|
def test_zero_experts_compute_with_zero_experts(self):
|
|
num_experts = 4
|
|
num_tokens = 3
|
|
hidden_size = 8
|
|
|
|
expert_indices = torch.tensor([[0, 1], [1, 2], [2, 3]], dtype=torch.int32)
|
|
expert_scales = torch.tensor([[0.5, 0.5], [0.3, 0.7], [0.4, 0.6]], dtype=torch.float32)
|
|
hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.float32)
|
|
|
|
result_indices, result_scales, result_hidden = zero_experts_compute(
|
|
expert_indices=expert_indices,
|
|
expert_scales=expert_scales,
|
|
num_experts=num_experts,
|
|
zero_expert_type="identity",
|
|
hidden_states=hidden_states,
|
|
)
|
|
|
|
assert result_indices.shape == expert_indices.shape
|
|
assert result_scales.shape == expert_scales.shape
|
|
assert result_hidden.shape == hidden_states.shape
|
|
|
|
def test_zero_experts_compute_normal_experts_masked(self):
|
|
num_experts = 4
|
|
num_tokens = 2
|
|
top_k = 2
|
|
hidden_size = 8
|
|
|
|
expert_indices = torch.tensor([[0, 5], [6, 7]], dtype=torch.int32)
|
|
expert_scales = torch.ones(num_tokens, top_k, dtype=torch.float32)
|
|
hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.float32)
|
|
|
|
result_indices, result_scales, _ = zero_experts_compute(
|
|
expert_indices=expert_indices,
|
|
expert_scales=expert_scales,
|
|
num_experts=num_experts,
|
|
zero_expert_type="identity",
|
|
hidden_states=hidden_states,
|
|
)
|
|
|
|
normal_expert_mask = expert_indices >= num_experts
|
|
for i in range(num_tokens):
|
|
for j in range(top_k):
|
|
if normal_expert_mask[i, j]:
|
|
assert result_scales[i, j] == 0.0
|
|
assert result_indices[i, j] == 0
|
|
|
|
def test_zero_experts_compute_output_sum(self):
|
|
num_experts = 2
|
|
num_tokens = 2
|
|
hidden_size = 4
|
|
|
|
expert_indices = torch.tensor([[0, 1], [0, 1]], dtype=torch.int32)
|
|
expert_scales = torch.tensor([[0.5, 0.5], [0.3, 0.7]], dtype=torch.float32)
|
|
hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.float32)
|
|
|
|
_, _, result_hidden = zero_experts_compute(
|
|
expert_indices=expert_indices,
|
|
expert_scales=expert_scales,
|
|
num_experts=num_experts,
|
|
zero_expert_type="identity",
|
|
hidden_states=hidden_states,
|
|
)
|
|
|
|
assert result_hidden.shape == (num_tokens, hidden_size)
|
|
|
|
def test_zero_experts_compute_all_zero_experts(self):
|
|
num_experts = 4
|
|
num_tokens = 2
|
|
top_k = 2
|
|
hidden_size = 8
|
|
|
|
expert_indices = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32)
|
|
expert_scales = torch.ones(num_tokens, top_k, dtype=torch.float32)
|
|
hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.float32)
|
|
|
|
result_indices, result_scales, result_hidden = zero_experts_compute(
|
|
expert_indices=expert_indices,
|
|
expert_scales=expert_scales,
|
|
num_experts=num_experts,
|
|
zero_expert_type="identity",
|
|
hidden_states=hidden_states,
|
|
)
|
|
|
|
assert torch.equal(result_indices, expert_indices)
|
|
assert torch.equal(result_scales, expert_scales)
|
|
assert result_hidden.shape == hidden_states.shape
|
|
|
|
def test_zero_experts_compute_mixed_experts(self):
|
|
num_experts = 3
|
|
num_tokens = 2
|
|
hidden_size = 8
|
|
|
|
expert_indices = torch.tensor([[0, 1, 4], [2, 5, 6]], dtype=torch.int32)
|
|
expert_scales = torch.tensor([[0.3, 0.3, 0.4], [0.2, 0.5, 0.3]], dtype=torch.float32)
|
|
hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.float32)
|
|
|
|
result_indices, result_scales, _ = zero_experts_compute(
|
|
expert_indices=expert_indices,
|
|
expert_scales=expert_scales,
|
|
num_experts=num_experts,
|
|
zero_expert_type="identity",
|
|
hidden_states=hidden_states,
|
|
)
|
|
|
|
assert result_indices.shape == expert_indices.shape
|
|
assert result_scales.shape == expert_scales.shape
|
|
|
|
def test_zero_experts_compute_identity_values_match_expected(self):
|
|
expert_indices = torch.tensor([[0, 2], [3, 1]], dtype=torch.int32)
|
|
expert_scales = torch.tensor([[0.25, 0.75], [0.60, 0.40]], dtype=torch.float32)
|
|
hidden_states = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
|
|
|
|
result_indices, result_scales, result_hidden = zero_experts_compute(
|
|
expert_indices=expert_indices,
|
|
expert_scales=expert_scales,
|
|
num_experts=2,
|
|
zero_expert_type="identity",
|
|
hidden_states=hidden_states,
|
|
)
|
|
|
|
expected_indices = torch.tensor([[0, 0], [0, 1]], dtype=torch.int32)
|
|
expected_scales = torch.tensor([[0.25, 0.0], [0.0, 0.40]], dtype=torch.float32)
|
|
expected_hidden = torch.tensor([[0.75, 1.50], [1.80, 2.40]], dtype=torch.float32)
|
|
|
|
assert torch.equal(result_indices, expected_indices)
|
|
torch.testing.assert_close(result_scales, expected_scales)
|
|
torch.testing.assert_close(result_hidden, expected_hidden)
|