141 lines
6.9 KiB
Python
141 lines
6.9 KiB
Python
from unittest.mock import MagicMock, Mock, patch
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from tests.ut.base import TestBase
|
|
from tests.ut.quantization.conftest_quantization import create_mock_ascend_config, create_mock_vllm_config
|
|
from vllm_ascend.quantization.methods.w4a4_mxfp4 import (
|
|
AscendW4A4MXFP4DynamicFusedMoEMethod,
|
|
AscendW4A4MXFP4DynamicLinearMethod,
|
|
)
|
|
|
|
|
|
class TestAscendW4A4MXFP4LinearMethod(TestBase):
|
|
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.ensure_mxfp4_linear_available")
|
|
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.get_current_vllm_config")
|
|
def setUp(self, mock_vllm, mock_ensure):
|
|
mock_vllm.return_value = create_mock_vllm_config()
|
|
mock_ensure.return_value = None
|
|
self.scheme = AscendW4A4MXFP4DynamicLinearMethod()
|
|
|
|
def test_get_weight_various_input_sizes(self):
|
|
for input_size in [64, 128, 256, 512]:
|
|
result = self.scheme.get_weight(input_size, 128, torch.bfloat16)
|
|
self.assertEqual(result["weight"].shape, (128, input_size // 2))
|
|
self.assertEqual(result["weight"].dtype, torch.uint8)
|
|
|
|
def test_get_pergroup_param_based_on_group_size(self):
|
|
group_sizes = [16, 32, 64]
|
|
for gs in group_sizes:
|
|
self.scheme.group_size = gs
|
|
result = self.scheme.get_pergroup_param(256, 128, torch.bfloat16)
|
|
self.assertEqual(result["weight_scale"].shape, (128, 256 // gs))
|
|
self.assertEqual(result["weight_scale"].dtype, torch.uint8)
|
|
|
|
def test_process_weights_after_loading_transposes(self):
|
|
layer = nn.Module()
|
|
layer.weight = nn.Parameter(torch.randint(0, 255, (128, 128), dtype=torch.uint8), requires_grad=False)
|
|
layer.weight_scale = nn.Parameter(torch.randint(0, 255, (128, 8), dtype=torch.uint8), requires_grad=False)
|
|
self.scheme.process_weights_after_loading(layer)
|
|
self.assertEqual(layer.weight.shape, (128, 128))
|
|
self.assertEqual(layer.weight_scale.shape[0], 4)
|
|
|
|
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.torch_npu")
|
|
def test_apply_3d_input(self, mock_npu):
|
|
mock_npu.npu_dynamic_mx_quant.return_value = (
|
|
torch.randint(0, 255, (32, 128), dtype=torch.uint8),
|
|
torch.randint(0, 255, (32, 4), dtype=torch.uint8),
|
|
)
|
|
mock_npu.npu_quant_matmul.return_value = torch.randn(32, 1, 128)
|
|
layer = MagicMock()
|
|
layer.weight = MagicMock(data=torch.randint(0, 255, (128, 128), dtype=torch.uint8))
|
|
layer.weight_scale = MagicMock(data=torch.randint(0, 255, (4, 128, 2), dtype=torch.uint8))
|
|
x = torch.randn(32, 1, 256, dtype=torch.bfloat16)
|
|
with patch.object(self.scheme, "group_size", 32):
|
|
output = self.scheme.apply(layer, x)
|
|
self.assertEqual(output.shape[0], 32)
|
|
|
|
|
|
class TestAscendW4A4MXFP4MoEMethod(TestBase):
|
|
num_experts = 8
|
|
hidden_size = 128
|
|
intermediate_size = 256
|
|
|
|
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.ensure_mxfp4_moe_available")
|
|
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.get_current_vllm_config")
|
|
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.get_ascend_config")
|
|
def setUp(self, mock_ascend, mock_vllm, mock_ensure):
|
|
mock_vllm.return_value = create_mock_vllm_config()
|
|
mock_ascend.return_value = create_mock_ascend_config()
|
|
mock_ensure.return_value = None
|
|
self.scheme = AscendW4A4MXFP4DynamicFusedMoEMethod()
|
|
|
|
def test_get_weight_static_method(self):
|
|
result = self.scheme.get_weight(self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16)
|
|
self.assertEqual(result["w13_weight"].dtype, torch.uint8)
|
|
self.assertEqual(result["w2_weight"].dtype, torch.uint8)
|
|
self.assertEqual(
|
|
result["w13_weight"].shape, (self.num_experts, 2 * self.intermediate_size, self.hidden_size // 2)
|
|
)
|
|
self.assertEqual(result["w2_weight"].shape, (self.num_experts, self.hidden_size, self.intermediate_size // 2))
|
|
|
|
def test_get_dynamic_quant_param_based_on_group_size(self):
|
|
group_sizes = [16, 32, 64]
|
|
for gs in group_sizes:
|
|
self.scheme.group_size = gs
|
|
result = self.scheme.get_dynamic_quant_param(
|
|
self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16
|
|
)
|
|
self.assertEqual(result["w13_weight_scale"].shape[2], self.hidden_size // gs)
|
|
self.assertEqual(result["w13_weight_scale"].dtype, torch.uint8)
|
|
self.assertEqual(result["w2_weight_scale"].dtype, torch.uint8)
|
|
|
|
def test_process_weights_transposes_weights(self):
|
|
layer = nn.Module()
|
|
layer.w13_weight = nn.Parameter(torch.randint(0, 255, (8, 256, 64), dtype=torch.uint8), requires_grad=False)
|
|
layer.w2_weight = nn.Parameter(torch.randint(0, 255, (8, 128, 128), dtype=torch.uint8), requires_grad=False)
|
|
layer.w13_weight_scale = nn.Parameter(
|
|
torch.randint(0, 255, (8, 256, 4), dtype=torch.uint8), requires_grad=False
|
|
)
|
|
layer.w2_weight_scale = nn.Parameter(torch.randint(0, 255, (8, 128, 8), dtype=torch.uint8), requires_grad=False)
|
|
self.scheme.process_weights_after_loading(layer)
|
|
self.assertEqual(layer.w13_weight.shape, (8, 64, 256))
|
|
self.assertEqual(layer.w13_weight_scale.shape, (8, 2, 256, 2))
|
|
|
|
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.torch_npu")
|
|
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4._EXTRA_CTX")
|
|
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.select_experts")
|
|
def test_apply_full_params(self, mock_select, mock_ctx, mock_npu):
|
|
tokens = 4
|
|
layer = nn.Module()
|
|
layer.w13_weight = nn.Parameter(torch.randint(0, 255, (8, 64, 256), dtype=torch.uint8), requires_grad=False)
|
|
layer.w2_weight = nn.Parameter(torch.randint(0, 255, (8, 128, 128), dtype=torch.uint8), requires_grad=False)
|
|
layer.w13_weight_scale = nn.Parameter(
|
|
torch.randint(0, 255, (8, 64, 128, 2), dtype=torch.uint8), requires_grad=False
|
|
)
|
|
layer.w2_weight_scale = nn.Parameter(
|
|
torch.randint(0, 255, (8, 128, 64, 2), dtype=torch.uint8), requires_grad=False
|
|
)
|
|
x = torch.randn(tokens, self.hidden_size, dtype=torch.bfloat16)
|
|
router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32)
|
|
topk_weights = torch.randn(tokens, 2)
|
|
topk_ids = torch.randint(0, self.num_experts, (tokens, 2))
|
|
mock_select.return_value = (topk_weights, topk_ids)
|
|
mock_comm = Mock()
|
|
mock_comm.fused_experts.return_value = torch.randn(tokens, self.hidden_size)
|
|
mock_ctx.moe_comm_method = mock_comm
|
|
mock_ctx.moe_comm_type = Mock()
|
|
self.scheme.apply(
|
|
layer,
|
|
x,
|
|
router_logits,
|
|
top_k=2,
|
|
renormalize=True,
|
|
num_experts=self.num_experts,
|
|
activation="silu",
|
|
pertoken_scale=torch.randn(tokens),
|
|
apply_router_weight_on_input=True,
|
|
)
|
|
mock_comm.fused_experts.assert_called_once()
|