from unittest.mock import MagicMock, Mock, patch import regex as re import torch from tests.ut.base import TestBase from tests.ut.quantization.conftest_quantization import identity from vllm_ascend.quantization.methods.w4a8 import AscendW4A8DynamicFusedMoEMethod, AscendW4A8DynamicLinearMethod from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD class TestAscendW4A8DynamicLinearMethod(TestBase): @patch("vllm_ascend.quantization.methods.w4a8.get_tensor_model_parallel_world_size") @patch("vllm_ascend.quantization.methods.w4a8.get_current_vllm_config") def setUp(self, mock_get_current_vllm_config, mock_get_tp_world_size): mock_get_tp_world_size.return_value = 1 mock_vllm_config = Mock() mock_vllm_config.quant_config = Mock(quant_description={"group_size": 256}) mock_vllm_config.scheduler_config = Mock( max_num_batched_tokens=2048, max_model_len=2048, enable_chunked_prefill=False ) mock_get_current_vllm_config.return_value = mock_vllm_config self.method = AscendW4A8DynamicLinearMethod() self.method.group_size = 8 def test_get_weight(self): weight = self.method.get_weight(8, 32, torch.bfloat16) self.assertEqual(weight["weight"].dtype, torch.int8) self.assertEqual(weight["weight"].shape, (32, 8)) # new quant version weight self.method.new_quant_version = True weight = self.method.get_weight(8, 32, torch.bfloat16) self.assertEqual(weight["weight"].dtype, torch.int8) self.assertEqual(weight["weight"].shape, (16, 8)) self.assertEqual(weight["_packed_dim"], 0) self.assertEqual(weight["_packed_factor"], 2) def test_get_pergroup_param(self): params = self.method.get_pergroup_param(8, 32, torch.bfloat16) self.assertEqual(params["weight_scale"].dtype, torch.bfloat16) self.assertEqual(params["weight_scale"].shape, (32, 1)) self.assertEqual(params["weight_offset"].dtype, torch.bfloat16) self.assertEqual(params["weight_offset"].shape, (32, 1)) self.assertEqual(params["weight_scale_second"].dtype, torch.bfloat16) self.assertEqual(params["weight_scale_second"].shape, (32, 1)) self.assertEqual(params["weight_offset_second"].dtype, torch.bfloat16) self.assertEqual(params["weight_offset_second"].shape, (32, 1)) # new quant version weight self.method.new_quant_version = True params = self.method.get_pergroup_param(8, 32, torch.bfloat16, layer_type="column") self.assertEqual(params["scale_bias"].dtype, torch.float32) self.assertEqual(params["scale_bias"].shape, (32, 1)) params = self.method.get_pergroup_param(8, 32, torch.bfloat16, layer_type="row") self.assertEqual(params["scale_bias"].dtype, torch.float32) self.assertEqual(params["scale_bias"].shape, (32, 16)) @patch("vllm_ascend.quantization.methods.w4a8.maybe_trans_nz") @patch("torch_npu.npu_convert_weight_to_int4pack") @patch("torch.Tensor.npu") @patch("torch_npu.npu_format_cast") def test_process_weights_after_loading( self, mock_format_cast, mock_npu, mock_npu_convert_weight, mock_maybe_trans_nz ): mock_npu.side_effect = lambda: torch.zeros((1, 32), dtype=torch.float32) mock_npu_convert_weight.return_value = torch.zeros((32, 4), dtype=torch.int32) mock_maybe_trans_nz.side_effect = identity # old quant version weight layer = torch.nn.Module() layer.weight = torch.nn.Parameter(torch.zeros((32, 8), dtype=torch.int8), requires_grad=False) layer.weight_scale = torch.nn.Parameter(torch.ones((32, 1), dtype=torch.float32), requires_grad=False) layer.weight_offset = torch.nn.Parameter(torch.empty_like(layer.weight_scale.data), requires_grad=False) layer.weight_scale_second = torch.nn.Parameter(torch.ones((32, 1), dtype=torch.float32), requires_grad=False) layer.weight_offset_second = torch.nn.Parameter( torch.empty_like(layer.weight_scale_second.data), requires_grad=False ) mock_format_cast.return_value = layer.weight.data.transpose(0, 1).contiguous() self.method.process_weights_after_loading(layer) self.assertTrue(hasattr(layer, "weight_scale_bias")) self.assertEqual(layer.weight_scale_bias.data.shape, (32,)) self.assertEqual(layer.weight_scale_bias.data.dtype, torch.float32) # new quant version weight self.method.new_quant_version = True new_layer = torch.nn.Module() new_layer.weight = torch.nn.Parameter(torch.zeros((16, 8), dtype=torch.int8), requires_grad=False) new_layer.weight_scale = torch.nn.Parameter(torch.ones((32, 1), dtype=torch.float32), requires_grad=False) new_layer.weight_offset = torch.nn.Parameter(torch.empty_like(new_layer.weight_scale.data), requires_grad=False) new_layer.weight_scale_second = torch.nn.Parameter( torch.ones((32, 1), dtype=torch.float32), requires_grad=False ) new_layer.weight_offset_second = torch.nn.Parameter( torch.empty_like(new_layer.weight_scale_second.data), requires_grad=False ) new_layer.scale_bias = torch.nn.Parameter(torch.zeros((32, 1), dtype=torch.float32), requires_grad=False) mock_format_cast.return_value = new_layer.weight.data.transpose(0, 1).contiguous() self.method.process_weights_after_loading(new_layer) self.assertEqual(new_layer.scale_bias.data.shape, (32,)) self.assertTrue(hasattr(new_layer, "weight_scale_second")) self.assertEqual(new_layer.weight_scale_second.data.shape, (1, 32)) @patch("torch_npu.npu_weight_quant_batchmatmul") def test_apply_basic(self, mock_matmul): layer = MagicMock() layer.weight = MagicMock(data=torch.randint(-8, 8, (256, 512), dtype=torch.int8)) layer.weight_scale_second = MagicMock(data=torch.randn(1, 512, dtype=torch.float32)) mock_matmul.return_value = torch.randn(32, 512) x = torch.randn(32, 256) self.method.apply(layer, x) mock_matmul.assert_called_once() @patch("vllm_ascend.quantization.methods.w4a8.maybe_trans_nz") def test_process_weights_after_loading_asserts_new_quant_packed_dim(self, mock_maybe_trans_nz): self.method.new_quant_version = True mock_maybe_trans_nz.side_effect = identity layer = torch.nn.Module() layer.weight = torch.nn.Parameter(torch.zeros((10, 16), dtype=torch.int8), requires_grad=False) layer.weight_scale = torch.nn.Parameter(torch.ones((20, 1), dtype=torch.float32), requires_grad=False) layer.weight_offset = torch.nn.Parameter(torch.empty_like(layer.weight_scale.data), requires_grad=False) layer.weight_scale_second = torch.nn.Parameter(torch.ones((20, 2), dtype=torch.float32), requires_grad=False) layer.weight_offset_second = torch.nn.Parameter( torch.empty_like(layer.weight_scale_second.data), requires_grad=False ) layer.scale_bias = torch.nn.Parameter(torch.zeros((20, 1), dtype=torch.float32), requires_grad=False) expected_message = "the last dim of weight needs to be divided by 4 but got shape torch.Size([16, 10])" with ( patch.object(self.method, "process_scale_second", return_value=(torch.ones((2, 20)), None)), self.assertRaisesRegex(AssertionError, re.escape(expected_message)), ): self.method.process_weights_after_loading(layer) class TestAscendW4A8DynamicLinearMethodWithNpu(TestBase): @patch("vllm_ascend.quantization.methods.w4a8.get_tensor_model_parallel_world_size") @patch("vllm_ascend.quantization.methods.w4a8.get_current_vllm_config") def setUp(self, mock_get_current_vllm_config, mock_get_tp_world_size): mock_get_tp_world_size.return_value = 1 mock_vllm_config = Mock() mock_vllm_config.quant_config = Mock(quant_description={"group_size": 64}) mock_get_current_vllm_config.return_value = mock_vllm_config self.method = AscendW4A8DynamicLinearMethod() def test_apply_with_npu(self): layer = torch.nn.Module() layer.weight = torch.nn.Parameter( torch.randint(-128, 127, (128, 32), dtype=torch.int32).npu(), requires_grad=False ) layer.weight_scale_second = torch.nn.Parameter( torch.randn(2, 256, dtype=torch.bfloat16).npu(), requires_grad=False ) x = torch.randn(32, 128, dtype=torch.bfloat16).npu() output = self.method.apply(layer, x) self.assertEqual(output.shape, (32, 256)) class TestAscendW4A8DynamicFusedMoEMethod(TestBase): experts = 8 input_size = 16 output_size = 56 group_size = 2 @patch("vllm_ascend.quantization.methods.w4a8.get_ascend_config") @patch("vllm_ascend.quantization.methods.w4a8.get_current_vllm_config") @patch("vllm_ascend.quantization.methods.w4a8.get_mc2_group") @patch("torch.distributed.get_rank", return_value=0) def setUp(self, mock_get_rank, mock_get_mc2_group, get_current_vllm_config, mock_get_ascend_config): # Mock ascend config mock_ascend_config = Mock() mock_ascend_config.eplb_config.dynamic_eplb = False mock_get_ascend_config.return_value = mock_ascend_config mock_vllm_config = Mock() mock_vllm_config.quant_config = Mock(quant_description={"group_size": self.group_size, "version": "0.0.0"}) mock_vllm_config.parallel_config = Mock(enable_expert_parallel=True) mock_vllm_config.scheduler_config = Mock( max_num_batched_tokens=2048, max_model_len=2048, enable_chunked_prefill=False ) get_current_vllm_config.return_value = mock_vllm_config self.quant_method = AscendW4A8DynamicFusedMoEMethod() def test_get_weight(self): # old quant version w4a8 weight param_dict = self.quant_method.get_weight(self.experts, self.input_size, self.output_size, torch.bfloat16) self.assertEqual(param_dict["w13_weight"].dtype, torch.int8) self.assertEqual(param_dict["w13_weight"].shape, (self.experts, 2 * self.input_size, self.output_size)) # new quant version weight self.quant_method.new_quant_version = True param_dict = self.quant_method.get_weight(self.experts, self.input_size, self.output_size, torch.bfloat16) self.assertEqual(param_dict["w13_weight"].dtype, torch.int8) self.assertEqual(param_dict["w13_weight"].shape, (self.experts, self.input_size, self.output_size)) def test_get_dynamic_quant_param(self): # old quant version weight param_dict = self.quant_method.get_dynamic_quant_param( self.experts, self.input_size, self.output_size, torch.bfloat16 ) self.assertEqual(param_dict["w13_weight_scale"].dtype, torch.float32) self.assertEqual(param_dict["w13_weight_scale"].shape, (self.experts, 2 * self.input_size, 1)) self.assertEqual(param_dict["w13_weight_scale_second"].dtype, torch.float32) self.assertEqual( param_dict["w13_weight_scale_second"].shape, (self.experts, 2 * self.input_size, self.output_size // self.group_size), ) self.assertEqual(param_dict["w2_weight_scale"].dtype, torch.float32) self.assertEqual(param_dict["w2_weight_scale"].shape, (self.experts, self.output_size, 1)) self.assertEqual(param_dict["w2_weight_scale_second"].dtype, torch.float32) self.assertEqual( param_dict["w2_weight_scale_second"].shape, (self.experts, self.output_size, self.input_size // self.group_size), ) # new quant version weight self.quant_method.new_quant_version = True param_dict = self.quant_method.get_dynamic_quant_param( self.experts, self.input_size, self.output_size, torch.bfloat16 ) self.assertEqual(param_dict["w2_scale_bias"].dtype, torch.float32) self.assertEqual( param_dict["w2_scale_bias"].shape, (self.experts, self.output_size, 16 // self.quant_method.tp_size) ) # per-channel weight self.quant_method.is_per_channel_weight = True param_dict = self.quant_method.get_dynamic_quant_param( self.experts, self.input_size, self.output_size, torch.bfloat16 ) pergroup_param = [ "w13_weight_scale_second", "w13_weight_offset_second", "w2_weight_scale_second", "w2_weight_offset_second", ] is_contains = any(key in param_dict for key in pergroup_param) self.assertFalse(is_contains) def build_layer(self, is_new_quant_version=True, is_per_channel_weight=False): layer = torch.nn.Module() if is_new_quant_version: layer.w13_weight = torch.nn.Parameter( torch.zeros((self.experts, self.input_size, self.output_size), dtype=torch.int8), requires_grad=False ) layer.w2_weight = torch.nn.Parameter( torch.zeros((self.experts, self.output_size // 2, self.input_size), dtype=torch.int8), requires_grad=False, ) w13_scale_bias = torch.zeros((self.experts, 2 * self.input_size, 1), dtype=torch.float32) layer.w13_scale_bias = torch.nn.Parameter(w13_scale_bias, requires_grad=False) w2_scale_bias = torch.zeros( (self.experts, self.output_size, 16 // self.quant_method.tp_size), dtype=torch.float32 ) layer.w2_scale_bias = torch.nn.Parameter(w2_scale_bias, requires_grad=False) else: layer.w13_weight = torch.nn.Parameter( torch.zeros((self.experts, 2 * self.input_size, self.output_size), dtype=torch.int8), requires_grad=False, ) layer.w2_weight = torch.nn.Parameter( torch.zeros((self.experts, self.output_size, self.input_size), dtype=torch.int8), requires_grad=False ) layer.w13_weight_scale = torch.nn.Parameter( torch.ones((self.experts, 2 * self.input_size, 1), dtype=torch.float32), requires_grad=False ) layer.w2_weight_scale = torch.nn.Parameter( torch.ones((self.experts, self.output_size, 1), dtype=torch.float32), requires_grad=False ) if not is_per_channel_weight: layer.w13_weight_scale_second = torch.nn.Parameter( torch.ones( (self.experts, 2 * self.input_size, self.output_size // self.group_size), dtype=torch.float32 ), requires_grad=False, ) layer.w13_weight_offset_second = torch.nn.Parameter( torch.empty_like(layer.w13_weight_scale_second.data), requires_grad=False ) layer.w2_weight_scale_second = torch.nn.Parameter( torch.ones((self.experts, self.output_size, self.input_size // self.group_size), dtype=torch.float32), requires_grad=False, ) layer.w2_weight_offset_second = torch.nn.Parameter( torch.empty_like(layer.w2_weight_scale_second.data), requires_grad=False ) return layer @patch("vllm_ascend.quantization.methods.w4a8.maybe_trans_nz") @patch("torch_npu.npu_format_cast") @patch("torch_npu.npu_quantize") @patch("torch.Tensor.npu", new=lambda self: self) def test_process_weights_after_loading(self, mock_npu_quantize, mock_npu_format_cast, mock_maybe_trans_nz): mock_npu_quantize.return_value = torch.Tensor() mock_npu_format_cast.side_effect = identity mock_maybe_trans_nz.side_effect = identity # old quant version weight layer = self.build_layer(is_new_quant_version=False) self.quant_method.process_weights_after_loading(layer) self.assertTrue(hasattr(layer, "w13_scale_bias")) self.assertEqual(layer.w13_scale_bias.data.shape, (self.experts, 2 * self.input_size)) self.assertEqual(layer.w13_scale_bias.data.dtype, torch.float32) self.assertTrue(hasattr(layer, "w2_scale_bias")) self.assertEqual(layer.w2_scale_bias.data.shape, (self.experts, self.output_size)) self.assertEqual(layer.w2_scale_bias.data.dtype, torch.float32) # new quant version weight self.quant_method.new_quant_version = True new_layer = self.build_layer(is_new_quant_version=True) self.quant_method.process_weights_after_loading(new_layer) self.assertEqual(new_layer.w13_scale_bias.data.shape, (self.experts, 2 * self.input_size)) self.assertEqual(new_layer.w2_scale_bias.data.shape, (self.experts, self.output_size)) self.assertFalse(hasattr(new_layer, "w13_weight_scale_second")) # per-channel weight self.quant_method.is_per_channel_weight = True per_channel_layer = self.build_layer(is_new_quant_version=True, is_per_channel_weight=True) self.quant_method.process_weights_after_loading(per_channel_layer) self.assertEqual(new_layer.w13_scale_bias.data.shape, (self.experts, 2 * self.input_size)) self.assertEqual(per_channel_layer.w13_weight_scale.data.shape, (self.experts, 2 * self.input_size)) def test_pack_to_int32_asserts_new_quant_packed_dim(self): self.quant_method.new_quant_version = True weight = torch.zeros((self.experts, self.output_size, 10), dtype=torch.int8) expected_message = f"the last dim of weight needs to be divided by 4 but got shape {weight.shape}" with self.assertRaisesRegex(AssertionError, re.escape(expected_message)): self.quant_method.pack_to_int32(weight) def test_get_weight_compressed_tensors(self): self.quant_method.quant_method = COMPRESSED_TENSORS_METHOD result = self.quant_method.get_weight(self.experts, self.input_size, self.output_size, torch.bfloat16) self.assertEqual(result["w13_weight"].dtype, torch.int8) def test_get_dynamic_quant_param_compressed_tensors(self): self.quant_method.quant_method = COMPRESSED_TENSORS_METHOD result = self.quant_method.get_dynamic_quant_param( self.experts, self.input_size, self.output_size, torch.bfloat16 ) self.assertIn("w13_weight_scale", result) self.assertIn("w2_weight_scale", result) self.assertEqual(result["w13_weight_scale"].dtype, torch.bfloat16) self.assertEqual(result["w2_weight_scale"].dtype, torch.bfloat16) @patch("vllm_ascend.quantization.methods.w4a8.maybe_trans_nz") @patch("torch_npu.npu_format_cast") @patch("torch_npu.npu_quantize") @patch("torch.Tensor.npu", new=lambda self: self) def test_process_weights_after_loading_compressed_tensors( self, mock_npu_quantize, mock_npu_format_cast, mock_maybe_trans_nz ): mock_npu_quantize.return_value = torch.Tensor() mock_npu_format_cast.side_effect = identity mock_maybe_trans_nz.side_effect = identity layer = self.build_layer(is_new_quant_version=False) self.quant_method.quant_method = COMPRESSED_TENSORS_METHOD self.quant_method.weight_strategy = "group" self.quant_method.process_weights_after_loading(layer) self.assertTrue(hasattr(layer, "w13_scale_bias")) self.assertEqual(layer.w13_scale_bias.data.shape, (self.experts, 2 * self.input_size)) self.assertEqual(layer.w13_scale_bias.data.dtype, torch.float32) self.quant_method.is_per_channel_weight = True self.quant_method.weight_strategy = "channel" per_channel_layer = self.build_layer(is_new_quant_version=False) self.quant_method.process_weights_after_loading(per_channel_layer) self.assertEqual(per_channel_layer.w13_weight_scale.data.shape, (self.experts, 2 * self.input_size)) self.assertEqual(per_channel_layer.w2_weight_scale.data.shape, (self.experts, 1, self.output_size)) @patch("vllm_ascend.quantization.methods.w4a8._EXTRA_CTX") @patch("vllm_ascend.quantization.methods.w4a8.select_experts") @patch("vllm_ascend.quantization.methods.w4a8.build_fused_experts_input") def test_apply_comprehensive(self, mock_build_input, mock_select, mock_ctx): tokens = 4 num_experts = self.experts hidden_size = self.output_size top_k = 2 layer = self.build_layer(is_new_quant_version=True, is_per_channel_weight=True) self.quant_method.is_per_channel_weight = True layer.swiglu_limit = 1000000 x = torch.randn(tokens, hidden_size, dtype=torch.bfloat16) router_logits = torch.randn(tokens, num_experts, dtype=torch.float32) topk_weights = torch.randn(tokens, top_k, dtype=torch.float32) topk_ids = torch.randint(0, num_experts, (tokens, top_k), dtype=torch.int64) expert_map = torch.randint(0, num_experts, (num_experts,), dtype=torch.int64) mc2_mask = torch.tensor([1, 0, 1, 0], dtype=torch.bool) pertoken_scale = torch.randn(tokens, dtype=torch.float32) log2phy = torch.randint(0, num_experts, (num_experts,), dtype=torch.int64) e_score_correction_bias = torch.randn(num_experts, dtype=torch.float32) mock_select.return_value = (topk_weights, topk_ids) mock_fused_input = Mock() mock_fused_input.hidden_states = x mock_fused_input.topk_weights = topk_weights mock_fused_input.topk_ids = topk_ids mock_fused_input.activation = "silu" mock_build_input.return_value = mock_fused_input mock_comm = Mock() expected_output = torch.randn(tokens, hidden_size, dtype=torch.bfloat16) mock_comm.fused_experts.return_value = expected_output mock_ctx.moe_comm_method = mock_comm output = self.quant_method.apply( layer=layer, x=x, router_logits=router_logits, top_k=top_k, renormalize=True, use_grouped_topk=False, num_experts=num_experts, expert_map=expert_map, scoring_func="softmax", routed_scaling_factor=1.0, e_score_correction_bias=e_score_correction_bias, is_prefill=True, enable_force_load_balance=False, log2phy=log2phy, global_redundant_expert_num=0, pertoken_scale=pertoken_scale, activation="silu", apply_router_weight_on_input=False, mc2_mask=mc2_mask, ) mock_select.assert_called_once() select_call_args = mock_select.call_args self.assertTrue(torch.equal(select_call_args.kwargs["hidden_states"], x)) self.assertEqual(select_call_args.kwargs["top_k"], top_k) self.assertEqual(select_call_args.kwargs["num_experts"], num_experts) mock_build_input.assert_called_once() build_kwargs = mock_build_input.call_args.kwargs self.assertTrue(torch.equal(build_kwargs["hidden_states"], x)) self.assertEqual(build_kwargs["quant_type"], self.quant_method.quant_type) self.assertTrue(build_kwargs["is_per_channel_weight"]) self.assertEqual(build_kwargs["activation"], "silu") self.assertEqual(build_kwargs["apply_router_weight_on_input"], False) mock_comm.fused_experts.assert_called_once() self.assertEqual(mock_comm.fused_experts.call_args.kwargs["fused_experts_input"], mock_fused_input) self.assertTrue(torch.equal(output, expected_output)) def test_apply_asserts_router_logits_expert_mismatch(self): layer = self.build_layer(is_new_quant_version=True, is_per_channel_weight=True) x = torch.randn(4, self.output_size, dtype=torch.bfloat16) router_logits = torch.randn(4, self.experts - 1, dtype=torch.float32) expected_message = ( "Number of global experts mismatch (excluding redundancy): " f"router_logits.shape[1]={self.experts - 1}, num_logical_experts={self.experts}" ) with self.assertRaisesRegex(AssertionError, re.escape(expected_message)): self.quant_method.apply( layer=layer, x=x, router_logits=router_logits, top_k=2, renormalize=True, num_experts=self.experts, )