71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.config import set_current_vllm_config
|
|
from vllm.model_executor.layers.layernorm import RMSNormGated
|
|
|
|
from vllm_ascend._310p.ops.layernorm import AscendRMSNormGated310
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def default_vllm_config():
|
|
mock_config = MagicMock()
|
|
mock_config.compilation_config.custom_ops = ["all"]
|
|
with set_current_vllm_config(mock_config):
|
|
yield mock_config
|
|
|
|
|
|
@patch("torch.nn.functional.silu", side_effect=lambda tensor: tensor + 1)
|
|
@patch("torch_npu.npu_rms_norm")
|
|
def test_rmsnorm_gated_310_forward_oot_uses_rmsnorm_activation_mul(mock_rms_norm, mock_silu):
|
|
layer = AscendRMSNormGated310(hidden_size=8, eps=1e-5, norm_before_gate=True)
|
|
x = torch.randn(2, 8, dtype=torch.float32)
|
|
z = torch.randn(2, 8, dtype=torch.float32)
|
|
normed = torch.randn(2, 8, dtype=torch.float32)
|
|
mock_rms_norm.return_value = (normed, None)
|
|
|
|
with patch.object(RMSNormGated, "forward_native", autospec=True) as mock_forward_native:
|
|
out = layer.forward_oot(x, z)
|
|
|
|
mock_forward_native.assert_not_called()
|
|
mock_rms_norm.assert_called_once()
|
|
rms_norm_args = mock_rms_norm.call_args.args
|
|
assert rms_norm_args[0] is x
|
|
assert rms_norm_args[1] is layer.weight
|
|
assert rms_norm_args[2] == layer.eps
|
|
mock_silu.assert_called_once_with(z)
|
|
assert torch.allclose(out, normed * (z + 1))
|
|
|
|
|
|
@patch("torch_npu.npu_rms_norm")
|
|
def test_rmsnorm_gated_310_forward_oot_uses_rmsnorm_without_gate(mock_rms_norm):
|
|
layer = AscendRMSNormGated310(hidden_size=8, eps=1e-5)
|
|
x = torch.randn(2, 8, dtype=torch.float32)
|
|
expected = torch.randn(2, 8, dtype=torch.float32)
|
|
mock_rms_norm.return_value = (expected, None)
|
|
|
|
with patch.object(RMSNormGated, "forward_native", autospec=True, return_value=expected) as mock_forward_native:
|
|
out = layer.forward_oot(x, None)
|
|
|
|
mock_forward_native.assert_not_called()
|
|
mock_rms_norm.assert_called_once()
|
|
rms_norm_args = mock_rms_norm.call_args.args
|
|
assert rms_norm_args[0] is x
|
|
assert rms_norm_args[1] is layer.weight
|
|
assert rms_norm_args[2] == layer.eps
|
|
assert out is expected
|
|
|
|
|
|
def test_rmsnorm_gated_310_forward_oot_keeps_native_for_group_norm():
|
|
layer = AscendRMSNormGated310(hidden_size=8, eps=1e-5, group_size=4)
|
|
x = torch.randn(2, 8, dtype=torch.float32)
|
|
z = torch.randn(2, 8, dtype=torch.float32)
|
|
expected = torch.randn(2, 8, dtype=torch.float32)
|
|
|
|
with patch.object(RMSNormGated, "forward_native", autospec=True, return_value=expected) as mock_forward_native:
|
|
out = layer.forward_oot(x, z)
|
|
|
|
mock_forward_native.assert_called_once_with(layer, x, z)
|
|
assert out is expected
|