From a1313ef474260f36371dbdcbb011c39cdfce7eb7 Mon Sep 17 00:00:00 2001 From: Wang Kunpeng <1289706727@qq.com> Date: Fri, 24 Apr 2026 16:14:45 +0800 Subject: [PATCH 1/4] [Test]Add quantization test case Signed-off-by: Wang Kunpeng <1289706727@qq.com> --- .../ut/quantization/conftest_quantization.py | 190 +++ .../quantization/{ => methods}/test_kv_c8.py | 1378 ++++++++--------- .../ut/quantization/methods/test_registry.py | 83 + .../quantization/{ => methods}/test_w4a16.py | 624 ++++---- .../test_w4a4_flatquant.py} | 35 +- .../methods/test_w4a4_laos_dynamic.py | 94 ++ .../quantization/methods/test_w4a4_mxfp4.py | 157 ++ .../test_w4a8.py} | 128 +- .../quantization/{ => methods}/test_w8a16.py | 52 +- .../quantization/methods/test_w8a8_dynamic.py | 339 ++++ .../quantization/methods/test_w8a8_mxfp8.py | 220 +++ .../quantization/methods/test_w8a8_pdmix.py | 164 ++ .../test_w8a8_static.py} | 124 +- .../test_compressed_tensors_config.py | 166 ++ tests/ut/quantization/test_method_adapters.py | 232 +++ .../ut/quantization/test_modelslim_config.py | 410 +++++ tests/ut/quantization/test_quant_parser.py | 123 ++ tests/ut/quantization/test_quant_type.py | 29 + .../{test_quant_utils.py => test_utils.py} | 36 +- tests/ut/quantization/test_w8a8_dynamic.py | 229 --- 20 files changed, 3505 insertions(+), 1308 deletions(-) create mode 100644 tests/ut/quantization/conftest_quantization.py rename tests/ut/quantization/{ => methods}/test_kv_c8.py (97%) create mode 100644 tests/ut/quantization/methods/test_registry.py rename tests/ut/quantization/{ => methods}/test_w4a16.py (95%) rename tests/ut/quantization/{test_w4a4_flatquant_dynamic.py => methods/test_w4a4_flatquant.py} (85%) create mode 100644 tests/ut/quantization/methods/test_w4a4_laos_dynamic.py create mode 100644 tests/ut/quantization/methods/test_w4a4_mxfp4.py rename tests/ut/quantization/{test_w4a8_dynamic.py => methods/test_w4a8.py} (73%) rename tests/ut/quantization/{ => methods}/test_w8a16.py (58%) create mode 100644 tests/ut/quantization/methods/test_w8a8_dynamic.py create mode 100644 tests/ut/quantization/methods/test_w8a8_mxfp8.py create mode 100644 tests/ut/quantization/methods/test_w8a8_pdmix.py rename tests/ut/quantization/{test_w8a8.py => methods/test_w8a8_static.py} (57%) create mode 100644 tests/ut/quantization/test_compressed_tensors_config.py create mode 100644 tests/ut/quantization/test_method_adapters.py create mode 100644 tests/ut/quantization/test_quant_parser.py create mode 100644 tests/ut/quantization/test_quant_type.py rename tests/ut/quantization/{test_quant_utils.py => test_utils.py} (86%) delete mode 100644 tests/ut/quantization/test_w8a8_dynamic.py diff --git a/tests/ut/quantization/conftest_quantization.py b/tests/ut/quantization/conftest_quantization.py new file mode 100644 index 000000000000..e5435de42038 --- /dev/null +++ b/tests/ut/quantization/conftest_quantization.py @@ -0,0 +1,190 @@ +from unittest.mock import Mock + +import torch +import torch.nn as nn + +FAKQUANT_CONFIG = { + "version": "1.0.0", + "model_quant_type": "W8A8_DYNAMIC", + "fa_quant_type": "FAKQuant", + "model.embed_tokens.weight": "FLOAT", + "model.layers.3.self_attn.fa_q.scale": "FAQuant", + "model.layers.3.self_attn.fa_k.scale": "FAQuant", + "model.layers.3.self_attn.fa_v.scale": "FAQuant", + "model.layers.3.self_attn.fa_q.offset": "FAQuant", + "model.layers.3.self_attn.fa_k.offset": "FAQuant", + "model.layers.3.self_attn.fa_v.offset": "FAQuant", +} + +W8A8_CONFIG = { + "version": "1.0.0", + "model_quant_type": "W8A8_DYNAMIC", + "model.embed_tokens.weight": "FLOAT", + "model.layers.0.self_attn.q_a_proj.weight": "W8A8", + "model.layers.0.mlp.gate_proj.weight": "W8A8_DYNAMIC", + "model.layers.0.mlp.up_proj.weight": "W8A8_DYNAMIC", + "model.layers.0.mlp.down_proj.weight": "W8A8_DYNAMIC", + "model.layers.3.mlp.experts.0.gate_proj.weight": "W8A8_DYNAMIC", + "model.layers.3.mlp.experts.0.up_proj.weight": "W8A8_DYNAMIC", + "model.layers.3.mlp.experts.0.down_proj.weight": "W8A8_DYNAMIC", + "model.layers.3.mlp.experts.1.gate_proj.weight": "W8A8_DYNAMIC", + "model.layers.3.mlp.experts.1.up_proj.weight": "W8A8_DYNAMIC", + "model.layers.3.mlp.experts.1.down_proj.weight": "W8A8_DYNAMIC", +} + +COMPRESSED_TENSORS_W8A8_CONFIG = { + "config_groups": { + "group_0": { + "format": "int-quantized", + "input_activations": { + "actorder": None, + "block_structure": None, + "dynamic": True, + "group_size": None, + "num_bits": 8, + "observer": None, + "observer_kwargs": {}, + "strategy": "token", + "symmetric": True, + "type": "int" + }, + "output_activations": None, + "targets": [ + "Linear" + ], + "weights": { + "actorder": None, + "block_structure": None, + "dynamic": False, + "group_size": None, + "num_bits": 8, + "observer": "minmax", + "observer_kwargs": {}, + "strategy": "channel", + "symmetric": True, + "type": "int" + } + } + }, + "format": "int-quantized", + "global_compression_ratio": None, + "ignore": [ + "lm_head" + ], + "kv_cache_scheme": None, + "quant_method": "compressed-tensors", + "quantization_status": "compressed", +} + +def identity(*args): + return args[0] + +def create_mock_vllm_config( + quant_description=None, + model_dtype=torch.bfloat16, + scheduler_config=None, + compilation_mode=None, + enforce_eager=True, + kv_transfer_config=None, + parallel_config=None, +): + if quant_description is None: + quant_description = {"group_size": 32} + + mock_config = Mock() + mock_config.quant_config = Mock(quant_description=quant_description) + mock_config.model_config = Mock( + dtype=model_dtype, + hf_config=Mock(model_type=None), + enforce_eager=enforce_eager, + ) + + if scheduler_config is None: + mock_config.scheduler_config = Mock( + max_num_batched_tokens=2048, + max_model_len=2048, + enable_chunked_prefill=False, + ) + else: + mock_config.scheduler_config = scheduler_config + + if compilation_mode is not None: + mock_config.compilation_config = Mock(mode=compilation_mode) + else: + mock_config.compilation_config = Mock() + + mock_config.kv_transfer_config = kv_transfer_config + + if parallel_config is None: + mock_config.parallel_config = Mock(enable_expert_parallel=True) + else: + mock_config.parallel_config = parallel_config + + return mock_config + + +def create_mock_ascend_config( + multistream_overlap_gate=False, + dynamic_eplb=False, + flashcomm2_oproj_tensor_parallel_size=1, +): + mock_config = Mock() + mock_config.multistream_overlap_gate = multistream_overlap_gate + mock_config.eplb_config = Mock(dynamic_eplb=dynamic_eplb) + mock_config.flashcomm2_oproj_tensor_parallel_size = flashcomm2_oproj_tensor_parallel_size + return mock_config + + +def create_moe_layer( + num_experts=8, + hidden_size=128, + intermediate_size=128, + weight_dtype=torch.int8, + params_dtype=torch.bfloat16, +): + layer = nn.Module() + layer.w13_weight = nn.Parameter( + torch.randint(-8, 8, (num_experts, 2 * intermediate_size, hidden_size), dtype=weight_dtype), + requires_grad=False, + ) + layer.w2_weight = nn.Parameter( + torch.randint(-8, 8, (num_experts, hidden_size, intermediate_size), dtype=weight_dtype), + requires_grad=False, + ) + layer.w13_weight_scale = nn.Parameter( + torch.ones((num_experts, 2 * intermediate_size, 1), dtype=params_dtype), requires_grad=False + ) + layer.w13_weight_offset = nn.Parameter( + torch.zeros((num_experts, 2 * intermediate_size, 1), dtype=params_dtype), requires_grad=False + ) + layer.w2_weight_scale = nn.Parameter( + torch.ones((num_experts, hidden_size, 1), dtype=params_dtype), requires_grad=False + ) + layer.w2_weight_offset = nn.Parameter( + torch.zeros((num_experts, hidden_size, 1), dtype=params_dtype), requires_grad=False + ) + return layer + + +def create_mxfp_moe_layer( + num_experts=8, + hidden_size=128, + intermediate_size=128, + group_size=32, + weight_dtype=torch.float8_e4m3fn, + scale_dtype=torch.uint8, +): + layer = nn.Module() + layer.w13_weight = nn.Parameter( + torch.randn(num_experts, 2 * intermediate_size, hidden_size).to(weight_dtype), requires_grad=False + ) + layer.w2_weight = nn.Parameter( + torch.randn(num_experts, hidden_size, intermediate_size).to(weight_dtype), requires_grad=False + ) + layer.w13_weight_scale = nn.Parameter( + torch.randint(0, 255, (num_experts, 2 * intermediate_size, hidden_size // group_size), dtype=scale_dtype), requires_grad=False + ) + layer.w2_weight_scale = nn.Parameter( + torch.randint(0, 255, (num_experts, hidden_size, intermediate_size // group_size), dtype=scale_dtype), requires_grad=False + ) + return layer diff --git a/tests/ut/quantization/test_kv_c8.py b/tests/ut/quantization/methods/test_kv_c8.py similarity index 97% rename from tests/ut/quantization/test_kv_c8.py rename to tests/ut/quantization/methods/test_kv_c8.py index 2173b999eb37..a98680b58b34 100644 --- a/tests/ut/quantization/test_kv_c8.py +++ b/tests/ut/quantization/methods/test_kv_c8.py @@ -1,689 +1,689 @@ -import unittest -import torch -import torch.nn as nn -from unittest.mock import MagicMock, Mock, patch - -from tests.ut.base import TestBase - - -class TestWeightLoader(unittest.TestCase): - """Test cases for weight_loader function in kv_c8.py""" - - def setUp(self): - """Set up test environment before each test""" - # Import the module under test - from vllm_ascend.quantization.methods.kv_c8 import _fa_quant_weight_loader as weight_loader - self.weight_loader = weight_loader - - # Mock distributed functions - self.tp_rank_patch = patch( - "vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_rank" - ) - self.tp_size_patch = patch( - "vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_world_size" - ) - self.mock_tp_rank = self.tp_rank_patch.start() - self.mock_tp_size = self.tp_size_patch.start() - - def tearDown(self): - """Clean up after each test""" - self.tp_rank_patch.stop() - self.tp_size_patch.stop() - - def test_weight_loader_single_element(self): - """Test weight_loader when both tensors contain a single element""" - # Create tensors with single element - param = torch.tensor([0.0]) - loaded_weight = torch.tensor([5.0]) - - # Call weight_loader - self.weight_loader(param, loaded_weight) - - # Verify the value was filled correctly - self.assertEqual(param.item(), 5.0) - self.assertEqual(param.dtype, torch.float32) - - def test_weight_loader_single_element_int(self): - """Test weight_loader with integer tensors""" - param = torch.tensor([0], dtype=torch.int32) - loaded_weight = torch.tensor([10], dtype=torch.int32) - - self.weight_loader(param, loaded_weight) - - self.assertEqual(param.item(), 10) - - def test_weight_loader_tp_sharding_first_rank(self): - """Test weight_loader with tensor parallelism sharding for first rank""" - # Configure mocks for rank 0 of 4 - self.mock_tp_rank.return_value = 0 - self.mock_tp_size.return_value = 4 - - # Create test tensors - param = torch.zeros(2, 5) # Target param shape (2,5) - loaded_weight = torch.ones(8, 5) # Full weight (8,5) - - # Mock narrow to track the call - with patch.object(loaded_weight, 'narrow', wraps=loaded_weight.narrow) as mock_narrow: - self.weight_loader(param, loaded_weight) - - # Verify narrow was called correctly: narrow(dim=0, start=0, length=2) - mock_narrow.assert_called_once_with(0, 0, 2) - - # Verify data was copied - self.assertTrue(torch.all(param == 1)) - - def test_weight_loader_tp_sharding_middle_rank(self): - """Test weight_loader with tensor parallelism sharding for middle rank""" - # Configure mocks for rank 2 of 4 - self.mock_tp_rank.return_value = 2 - self.mock_tp_size.return_value = 4 - - param = torch.zeros(2, 5) - loaded_weight = torch.ones(8, 5) - - with patch.object(loaded_weight, 'narrow', wraps=loaded_weight.narrow) as mock_narrow: - self.weight_loader(param, loaded_weight) - - # Verify narrow was called correctly: start = shard_size * rank = 2 * 2 = 4 - mock_narrow.assert_called_once_with(0, 4, 2) - - self.assertTrue(torch.all(param == 1)) - - def test_weight_loader_tp_sharding_last_rank(self): - """Test weight_loader with tensor parallelism sharding for last rank""" - # Configure mocks for rank 3 of 4 - self.mock_tp_rank.return_value = 3 - self.mock_tp_size.return_value = 4 - - param = torch.zeros(2, 5) - loaded_weight = torch.ones(8, 5) - - with patch.object(loaded_weight, 'narrow', wraps=loaded_weight.narrow) as mock_narrow: - self.weight_loader(param, loaded_weight) - - # Verify narrow was called correctly: start = 2 * 3 = 6 - mock_narrow.assert_called_once_with(0, 6, 2) - - self.assertTrue(torch.all(param == 1)) - - def test_weight_loader_shape_mismatch(self): - """Test weight_loader raises assertion error on shape mismatch""" - self.mock_tp_rank.return_value = 0 - self.mock_tp_size.return_value = 2 - - param = torch.zeros(2, 3) - loaded_weight = torch.ones(4, 4) # Different shape after sharding - - # Mock narrow to return tensor with wrong shape - with patch.object(loaded_weight, 'narrow', return_value=torch.ones(2, 4)): - with self.assertRaises(AssertionError) as context: - self.weight_loader(param, loaded_weight) - - # Verify error message contains expected information - self.assertIn("Attempted to load weight", str(context.exception)) - self.assertIn("into parameter", str(context.exception)) - - def test_weight_loader_with_different_dtypes(self): - """Test weight_loader handles different dtypes correctly""" - self.mock_tp_rank.return_value = 0 - self.mock_tp_size.return_value = 1 # No sharding - - param = torch.zeros(2, 3, dtype=torch.float32) - loaded_weight = torch.ones(2, 3, dtype=torch.float16) - - self.weight_loader(param, loaded_weight) - - # Verify data was copied and converted - self.assertTrue(torch.all(param == 1)) - self.assertEqual(param.dtype, torch.float32) - - -class TestAscendFAQuantAttentionMethodInit(unittest.TestCase): - """Test cases for AscendFAQuantAttentionMethod initialization""" - - def setUp(self): - """Set up test environment""" - # Mock vllm_config - self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config") - self.mock_get_config = self.config_patch.start() - - # Create mock config with attributes - self.mock_config = Mock() - self.mock_hf_config = Mock() - self.mock_hf_config.kv_lora_rank = 128 - self.mock_hf_config.qk_rope_head_dim = 64 - self.mock_config.model_config.hf_config = self.mock_hf_config - self.mock_get_config.return_value = self.mock_config - - # Import the class after patching - from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod - self.method_class = AscendFAQuantAttentionMethod - - def tearDown(self): - """Clean up after each test""" - self.config_patch.stop() - - def test_init_with_full_config(self): - """Test initialization when config has all attributes""" - method = self.method_class() - - self.assertTrue(method.transpose_weight) - self.assertFalse(method.printFlag) - self.assertEqual(method.kv_lora_rank, 128) - self.assertEqual(method.qk_rope_head_dim, 64) - - def test_init_without_kv_lora_rank(self): - """Test initialization when config lacks kv_lora_rank""" - delattr(self.mock_hf_config, "kv_lora_rank") - - method = self.method_class() - - self.assertEqual(method.kv_lora_rank, 0) - self.assertEqual(method.qk_rope_head_dim, 64) - - def test_init_without_qk_rope_head_dim(self): - """Test initialization when config lacks qk_rope_head_dim""" - delattr(self.mock_hf_config, "qk_rope_head_dim") - - method = self.method_class() - - self.assertEqual(method.kv_lora_rank, 128) - self.assertEqual(method.qk_rope_head_dim, 0) - - def test_init_without_both_attributes(self): - """Test initialization when config lacks both attributes""" - delattr(self.mock_hf_config, "kv_lora_rank") - delattr(self.mock_hf_config, "qk_rope_head_dim") - - method = self.method_class() - - self.assertEqual(method.kv_lora_rank, 0) - self.assertEqual(method.qk_rope_head_dim, 0) - - -class TestAscendFAQuantAttentionMethodCreateWeights(unittest.TestCase): - """Test cases for create_weights method""" - - def setUp(self): - """Set up test environment""" - # Mock vllm_config - self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config") - self.mock_get_config = self.config_patch.start() - - self.mock_config = Mock() - self.mock_hf_config = Mock() - self.mock_hf_config.kv_lora_rank = 128 - self.mock_hf_config.qk_rope_head_dim = 64 - self.mock_config.model_config.hf_config = self.mock_hf_config - self.mock_get_config.return_value = self.mock_config - - # Import the class - from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod - self.method_class = AscendFAQuantAttentionMethod - - # Mock torch functions - self.default_dtype_patch = patch("torch.get_default_dtype", return_value=torch.float32) - self.mock_default_dtype = self.default_dtype_patch.start() - - # Create a real nn.Module for testing - self.layer = nn.Module() - self.layer.num_heads = 32 - self.layer.num_kv_heads = 1 - - def tearDown(self): - """Clean up after each test""" - self.config_patch.stop() - self.default_dtype_patch.stop() - - def test_create_weights_adds_submodules(self): - """Test that create_weights adds fa_q, fa_k, fa_v submodules""" - method = self.method_class() - - with patch("torch.empty") as mock_empty: - mock_empty.return_value = torch.zeros(1, 1) - - method.create_weights(self.layer) - - # Verify submodules were added - self.assertTrue(hasattr(self.layer, "fa_q")) - self.assertTrue(hasattr(self.layer, "fa_k")) - self.assertTrue(hasattr(self.layer, "fa_v")) - - # Verify they are instances of nn.Module - self.assertIsInstance(self.layer.fa_q, nn.Module) - self.assertIsInstance(self.layer.fa_k, nn.Module) - self.assertIsInstance(self.layer.fa_v, nn.Module) - - def test_create_weights_creates_correct_tensors(self): - """Test that create_weights creates tensors with correct shapes and dtypes""" - method = self.method_class() - - # Track torch.empty calls - empty_calls = [] - - def mock_empty(size, dtype=None): - empty_calls.append((size, dtype)) - return torch.zeros(size, dtype=dtype if dtype else torch.float32) - - with patch("torch.empty", side_effect=mock_empty): - method.create_weights(self.layer) - - # Verify tensor creations - expected_calls = [ - ((32, 1), torch.float32), # fa_q.scale - ((1, 1), torch.float32), # fa_k.scale - ((1, 1), torch.float32), # fa_v.scale - ((32, 1), torch.int8), # fa_q.offset - ((1, 1), torch.int8), # fa_k.offset - ((1, 1), torch.int8), # fa_v.offset - ] - - # Compare without considering order - self.assertEqual(len(empty_calls), len(expected_calls)) - for call in expected_calls: - self.assertIn(call, empty_calls) - - def test_create_weights_registers_parameters(self): - """Test that create_weights registers parameters with correct attributes""" - method = self.method_class() - - # Create real tensors for testing - def create_tensor(*args, **kwargs): - size = args[0] if args else kwargs.get('size', (1,)) - dtype = kwargs.get('dtype', torch.float32) - return torch.zeros(*size, dtype=dtype) - - with patch("torch.empty", side_effect=create_tensor): - method.create_weights(self.layer) - - # Import weight_loader for comparison - from vllm_ascend.quantization.methods.kv_c8 import _fa_quant_weight_loader as weight_loader - - # Verify each parameter exists and has weight_loader - self.assertTrue(hasattr(self.layer.fa_q, "scale")) - self.assertTrue(hasattr(self.layer.fa_q.scale, "weight_loader")) - self.assertEqual(self.layer.fa_q.scale.weight_loader, weight_loader) - self.assertFalse(self.layer.fa_q.scale.requires_grad) - - self.assertTrue(hasattr(self.layer.fa_k, "scale")) - self.assertTrue(hasattr(self.layer.fa_k.scale, "weight_loader")) - - self.assertTrue(hasattr(self.layer.fa_v, "scale")) - self.assertTrue(hasattr(self.layer.fa_v.scale, "weight_loader")) - - self.assertTrue(hasattr(self.layer.fa_q, "offset")) - self.assertTrue(hasattr(self.layer.fa_q.offset, "weight_loader")) - self.assertEqual(self.layer.fa_q.offset.dtype, torch.int8) - - -class TestAscendFAQuantAttentionMethodProcessWeights(unittest.TestCase): - """Test cases for process_weights_after_loading method""" - - def setUp(self): - """Set up test environment""" - # Mock vllm_config - self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config") - self.mock_get_config = self.config_patch.start() - - self.mock_config = Mock() - self.mock_hf_config = Mock() - self.mock_hf_config.kv_lora_rank = 64 - self.mock_hf_config.qk_rope_head_dim = 32 - self.mock_config.model_config.hf_config = self.mock_hf_config - self.mock_get_config.return_value = self.mock_config - - # Import the class - from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod - self.method_class = AscendFAQuantAttentionMethod - - # Create method instance with real layer - self.method = self.method_class() - - # Create a real nn.Module for testing - self.layer = nn.Module() - - # Create real tensors for fa_k - self.fa_k_scale = torch.tensor([[2.0, 3.0, 4.0]], dtype=torch.float16) # Shape (1,3) - self.fa_k_offset = torch.tensor([[1, 2, 3]], dtype=torch.int8) # Shape (1,3) - - # Create fa_k module with parameters - self.layer.fa_k = nn.Module() - self.layer.fa_k.scale = nn.Parameter(self.fa_k_scale, requires_grad=False) - self.layer.fa_k.offset = nn.Parameter(self.fa_k_offset, requires_grad=False) - - def tearDown(self): - """Clean up after each test""" - self.config_patch.stop() - - def test_process_weights_with_single_value_scale(self): - """Test process_weights with single value scale""" - # Create new layer with single value scale - layer = nn.Module() - layer.fa_k = nn.Module() - layer.fa_k.scale = nn.Parameter(torch.tensor([[2.0]], dtype=torch.float16), requires_grad=False) - layer.fa_k.offset = nn.Parameter(torch.tensor([[1]], dtype=torch.int8), requires_grad=False) - - self.method.kv_lora_rank = 4 - self.method.process_weights_after_loading(layer) - - self.assertEqual(layer.quant_kscale.shape, (1, 4)) - self.assertEqual(layer.quant_kscale.dtype, torch.float32) - - -class TestIntegration(unittest.TestCase): - """Integration tests for the complete kv_c8 functionality""" - - def setUp(self): - """Set up test environment""" - # Mock vllm_config - self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config") - self.mock_get_config = self.config_patch.start() - - self.mock_config = Mock() - self.mock_hf_config = Mock() - self.mock_hf_config.kv_lora_rank = 64 - self.mock_hf_config.qk_rope_head_dim = 32 - self.mock_config.model_config.hf_config = self.mock_hf_config - self.mock_get_config.return_value = self.mock_config - - # Mock distributed functions - self.tp_rank_patch = patch( - "vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_rank" - ) - self.tp_size_patch = patch( - "vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_world_size" - ) - self.mock_tp_rank = self.tp_rank_patch.start() - self.mock_tp_size = self.tp_size_patch.start() - - def tearDown(self): - """Clean up after each test""" - self.config_patch.stop() - self.tp_rank_patch.stop() - self.tp_size_patch.stop() - - def test_complete_workflow(self): - """Test complete workflow from weight creation to processing""" - from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod - - # Create method instance - method = AscendFAQuantAttentionMethod() - - # Create real layer - layer = nn.Module() - layer.num_heads = 32 - layer.num_kv_heads = 1 - - # Step 1: Create weights - method.create_weights(layer) - - # Verify weights were created with correct structure - self.assertTrue(hasattr(layer, "fa_q")) - self.assertTrue(hasattr(layer, "fa_k")) - self.assertTrue(hasattr(layer, "fa_v")) - - self.assertTrue(hasattr(layer.fa_q, "scale")) - self.assertTrue(hasattr(layer.fa_q, "offset")) - self.assertTrue(hasattr(layer.fa_k, "scale")) - self.assertTrue(hasattr(layer.fa_k, "offset")) - self.assertTrue(hasattr(layer.fa_v, "scale")) - self.assertTrue(hasattr(layer.fa_v, "offset")) - - # Step 2: Simulate weight loading - self.mock_tp_rank.return_value = 0 - self.mock_tp_size.return_value = 1 - - # Create dummy weights - q_scale = torch.randn(32, 1) - k_scale = torch.randn(1, 1) - v_scale = torch.randn(1, 1) - q_offset = torch.randint(-128, 127, (32, 1), dtype=torch.int8) - k_offset = torch.randint(-128, 127, (1, 1), dtype=torch.int8) - v_offset = torch.randint(-128, 127, (1, 1), dtype=torch.int8) - - # Load weights using weight_loader - from vllm_ascend.quantization.methods.kv_c8 import _fa_quant_weight_loader as weight_loader - - with torch.no_grad(): - weight_loader(layer.fa_q.scale, q_scale) - weight_loader(layer.fa_k.scale, k_scale) - weight_loader(layer.fa_v.scale, v_scale) - weight_loader(layer.fa_q.offset, q_offset) - weight_loader(layer.fa_k.offset, k_offset) - weight_loader(layer.fa_v.offset, v_offset) - - # Verify weights were loaded correctly - self.assertTrue(torch.all(layer.fa_q.scale == q_scale)) - self.assertTrue(torch.all(layer.fa_k.scale == k_scale)) - self.assertTrue(torch.all(layer.fa_v.scale == v_scale)) - - # Step 3: Process after loading - method.process_weights_after_loading(layer) - - # Verify processed parameters - self.assertTrue(hasattr(layer, "fak_descale")) - self.assertTrue(hasattr(layer, "fak_offset")) - self.assertTrue(hasattr(layer, "quant_kscale")) - - -class TestC8KVScaleWeightLoader(TestBase): - """Tests for _c8_kv_scale_weight_loader in kv_c8.py.""" - - def setUp(self): - from vllm_ascend.quantization.methods.kv_c8 import _c8_kv_scale_weight_loader - self.loader = _c8_kv_scale_weight_loader - - def test_shape_match_copies_value(self): - param = nn.Parameter(torch.ones(4, dtype=torch.float32), requires_grad=False) - loaded = torch.tensor([1.0, 2.0, 3.0, 4.0]) - self.loader(param, loaded) - self.assertTrue(torch.allclose(param.data, loaded.float())) - - def test_shape_mismatch_resizes_param(self): - param = nn.Parameter(torch.ones(1, dtype=torch.float32), requires_grad=False) - loaded = torch.arange(8, dtype=torch.float32) - self.loader(param, loaded) - self.assertEqual(param.data.shape, (8,)) - self.assertTrue(torch.allclose(param.data, loaded)) - - def test_squeeze_before_compare(self): - param = nn.Parameter(torch.ones(4, dtype=torch.float32), requires_grad=False) - loaded = torch.arange(4, dtype=torch.float32).unsqueeze(0) # shape [1, 4] - self.loader(param, loaded) - self.assertEqual(param.data.shape, (4,)) - - def test_dtype_preserved_as_param_dtype(self): - param = nn.Parameter(torch.ones(4, dtype=torch.float32), requires_grad=False) - loaded = torch.arange(4, dtype=torch.float16) - self.loader(param, loaded) - self.assertEqual(param.data.dtype, torch.float32) - - -class TestAscendC8KVCacheAttentionMethod(TestBase): - """Tests for AscendC8KVCacheAttentionMethod in kv_c8.py.""" - - def _make_method(self): - from vllm_ascend.quantization.methods.kv_c8 import AscendC8KVCacheAttentionMethod - return AscendC8KVCacheAttentionMethod(quant_description={}, prefix="model.layers.0.self_attn.attn") - - def _make_layer_with_impl(self): - layer = nn.Module() - layer.impl = MagicMock() - return layer - - def test_create_weights_sets_kv_cache_torch_dtype(self): - method = self._make_method() - layer = self._make_layer_with_impl() - method.create_weights(layer) - self.assertEqual(layer.kv_cache_torch_dtype, torch.int8) - - def test_create_weights_registers_scale_offset_params(self): - method = self._make_method() - layer = self._make_layer_with_impl() - method.create_weights(layer) - self.assertIsInstance(layer.k_cache_scale, nn.Parameter) - self.assertIsInstance(layer.k_cache_offset, nn.Parameter) - self.assertIsInstance(layer.v_cache_scale, nn.Parameter) - self.assertIsInstance(layer.v_cache_offset, nn.Parameter) - self.assertFalse(layer.k_cache_scale.requires_grad) - self.assertFalse(layer.v_cache_offset.requires_grad) - - def test_create_weights_initial_values(self): - method = self._make_method() - layer = self._make_layer_with_impl() - method.create_weights(layer) - self.assertEqual(layer.k_cache_scale.data.item(), 1.0) - self.assertEqual(layer.v_cache_scale.data.item(), 1.0) - self.assertEqual(layer.k_cache_offset.data.item(), 0.0) - self.assertEqual(layer.v_cache_offset.data.item(), 0.0) - - def test_create_weights_assigns_weight_loader(self): - from vllm_ascend.quantization.methods.kv_c8 import _c8_kv_scale_weight_loader - method = self._make_method() - layer = self._make_layer_with_impl() - method.create_weights(layer) - self.assertIs(layer.k_cache_scale.weight_loader, _c8_kv_scale_weight_loader) - self.assertIs(layer.v_cache_scale.weight_loader, _c8_kv_scale_weight_loader) - self.assertIs(layer.k_cache_offset.weight_loader, _c8_kv_scale_weight_loader) - self.assertIs(layer.v_cache_offset.weight_loader, _c8_kv_scale_weight_loader) - - def test_process_weights_after_loading_flattens(self): - method = self._make_method() - layer = nn.Module() - layer.k_cache_scale = nn.Parameter(torch.ones(2, 4), requires_grad=False) - layer.k_cache_offset = nn.Parameter(torch.zeros(2, 4), requires_grad=False) - layer.v_cache_scale = nn.Parameter(torch.ones(2, 4), requires_grad=False) - layer.v_cache_offset = nn.Parameter(torch.zeros(2, 4), requires_grad=False) - method.process_weights_after_loading(layer) - self.assertEqual(layer.k_cache_scale.data.dim(), 1) - self.assertEqual(layer.k_cache_scale.data.shape[0], 8) - self.assertEqual(layer.v_cache_offset.data.dim(), 1) - - def test_apply_raises_runtime_error(self): - method = self._make_method() - layer = MagicMock() - with self.assertRaises(RuntimeError): - method.apply(layer, MagicMock(), MagicMock(), MagicMock(), None, None, None, None, None) - - -class TestAscendC8AttentionBackendImplScales(TestBase): - """Tests for AscendC8AttentionBackendImpl scale helpers.""" - - def _make_impl(self, num_kv_heads=4, head_size=8): - from vllm_ascend.attention.attention_v1 import AscendC8AttentionBackendImpl - impl = object.__new__(AscendC8AttentionBackendImpl) - impl.num_heads = num_kv_heads - impl.num_kv_heads = num_kv_heads - impl.head_size = head_size - impl.scale = 1.0 - impl.key_cache = None - impl.value_cache = None - return impl - - def _make_layer(self, num_kv_heads=4, head_size=8): - layer = nn.Module() - layer.k_cache_scale = nn.Parameter( - torch.ones(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False - ) - layer.k_cache_offset = nn.Parameter( - torch.zeros(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False - ) - layer.v_cache_scale = nn.Parameter( - torch.ones(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False - ) - layer.v_cache_offset = nn.Parameter( - torch.zeros(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False - ) - return layer - - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) - def test_prepare_c8_scales_runs_once(self, mock_tp_size, mock_tp_rank): - impl = self._make_impl() - layer = self._make_layer() - impl._prepare_c8_scales(layer, torch.device("cpu")) - self.assertTrue(hasattr(layer, "_c8_scales_prepared")) - self.assertTrue(layer._c8_scales_prepared) - - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) - def test_prepare_c8_scales_idempotent(self, mock_tp_size, mock_tp_rank): - impl = self._make_impl() - layer = self._make_layer() - impl._prepare_c8_scales(layer, torch.device("cpu")) - k_scale_after_first = layer._c8_k_scale.clone() - layer.k_cache_scale.data = torch.ones(32, dtype=torch.float32) * 99 - impl._prepare_c8_scales(layer, torch.device("cpu")) - self.assertTrue(torch.allclose(layer._c8_k_scale, k_scale_after_first)) - - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) - def test_prepare_c8_scales_creates_bnsd_shape(self, mock_tp_size, mock_tp_rank): - num_kv_heads, head_size = 4, 8 - impl = self._make_impl(num_kv_heads, head_size) - layer = self._make_layer(num_kv_heads, head_size) - impl._prepare_c8_scales(layer, torch.device("cpu")) - self.assertEqual(layer._c8_k_aq_scale.shape, (1, num_kv_heads, 1, head_size)) - self.assertEqual(layer._c8_v_aq_scale.shape, (1, num_kv_heads, 1, head_size)) - self.assertEqual(layer._c8_k_aq_scale.dtype, torch.bfloat16) - - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) - def test_quantize_kv_to_int8_output_dtype(self, mock_tp_size, mock_tp_rank): - num_kv_heads, head_size = 4, 8 - impl = self._make_impl(num_kv_heads, head_size) - layer = self._make_layer(num_kv_heads, head_size) - impl._prepare_c8_scales(layer, torch.device("cpu")) - num_tokens = 6 - key = torch.zeros(num_tokens, num_kv_heads, head_size, dtype=torch.bfloat16) - value = torch.zeros(num_tokens, num_kv_heads, head_size, dtype=torch.bfloat16) - key_q, value_q = impl._quantize_kv_to_int8(key, value, layer, num_tokens) - self.assertEqual(key_q.dtype, torch.int8) - self.assertEqual(value_q.dtype, torch.int8) - self.assertEqual(key_q.shape, key.shape) - - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) - def test_quantize_kv_to_int8_formula(self, mock_tp_size, mock_tp_rank): - """With scale=2.0, offset=0: q = round(x / 2).""" - num_kv_heads, head_size = 1, 4 - impl = self._make_impl(num_kv_heads, head_size) - layer = nn.Module() - scale_val = torch.full((num_kv_heads * head_size,), 2.0, dtype=torch.float32) - layer.k_cache_scale = nn.Parameter(scale_val.clone(), requires_grad=False) - layer.k_cache_offset = nn.Parameter(torch.zeros(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False) - layer.v_cache_scale = nn.Parameter(scale_val.clone(), requires_grad=False) - layer.v_cache_offset = nn.Parameter(torch.zeros(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False) - impl._prepare_c8_scales(layer, torch.device("cpu")) - key = torch.full((1, num_kv_heads, head_size), 4.0, dtype=torch.bfloat16) - value = torch.full((1, num_kv_heads, head_size), 4.0, dtype=torch.bfloat16) - key_q, _ = impl._quantize_kv_to_int8(key, value, layer, 1) - self.assertTrue(torch.all(key_q[0] == 2)) - - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) - @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) - def test_dequant_paged_kv_to_dense_round_trip(self, mock_tp_size, mock_tp_rank): - """With scale=1, offset=0: dequant(int8) == float(int8).""" - num_kv_heads, head_size = 2, 4 - block_size = 32 - num_blocks = 2 - H = num_kv_heads * head_size - impl = self._make_impl(num_kv_heads, head_size) - layer = self._make_layer(num_kv_heads, head_size) - impl._prepare_c8_scales(layer, torch.device("cpu")) - - key_int8 = torch.randint(-10, 10, (num_blocks, block_size, H), dtype=torch.int8) - value_int8 = torch.randint(-10, 10, (num_blocks, block_size, H), dtype=torch.int8) - seq_lens = [32, 32] - block_table = torch.tensor([[0], [1]], dtype=torch.long) - - dense_k, dense_v = impl._dequant_paged_kv_to_dense( - key_int8, value_int8, block_table, seq_lens, torch.float32, layer - ) - expected_k = key_int8.view(-1, num_kv_heads, head_size).float() - self.assertEqual(dense_k.shape, (64, num_kv_heads, head_size)) - self.assertTrue(torch.allclose(dense_k, expected_k)) - - -if __name__ == "__main__": - unittest.main(verbosity=2) +import unittest +import torch +import torch.nn as nn +from unittest.mock import MagicMock, Mock, patch + +from tests.ut.base import TestBase + + +class TestWeightLoader(unittest.TestCase): + """Test cases for weight_loader function in kv_c8.py""" + + def setUp(self): + """Set up test environment before each test""" + # Import the module under test + from vllm_ascend.quantization.methods.kv_c8 import _fa_quant_weight_loader as weight_loader + self.weight_loader = weight_loader + + # Mock distributed functions + self.tp_rank_patch = patch( + "vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_rank" + ) + self.tp_size_patch = patch( + "vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_world_size" + ) + self.mock_tp_rank = self.tp_rank_patch.start() + self.mock_tp_size = self.tp_size_patch.start() + + def tearDown(self): + """Clean up after each test""" + self.tp_rank_patch.stop() + self.tp_size_patch.stop() + + def test_weight_loader_single_element(self): + """Test weight_loader when both tensors contain a single element""" + # Create tensors with single element + param = torch.tensor([0.0]) + loaded_weight = torch.tensor([5.0]) + + # Call weight_loader + self.weight_loader(param, loaded_weight) + + # Verify the value was filled correctly + self.assertEqual(param.item(), 5.0) + self.assertEqual(param.dtype, torch.float32) + + def test_weight_loader_single_element_int(self): + """Test weight_loader with integer tensors""" + param = torch.tensor([0], dtype=torch.int32) + loaded_weight = torch.tensor([10], dtype=torch.int32) + + self.weight_loader(param, loaded_weight) + + self.assertEqual(param.item(), 10) + + def test_weight_loader_tp_sharding_first_rank(self): + """Test weight_loader with tensor parallelism sharding for first rank""" + # Configure mocks for rank 0 of 4 + self.mock_tp_rank.return_value = 0 + self.mock_tp_size.return_value = 4 + + # Create test tensors + param = torch.zeros(2, 5) # Target param shape (2,5) + loaded_weight = torch.ones(8, 5) # Full weight (8,5) + + # Mock narrow to track the call + with patch.object(loaded_weight, 'narrow', wraps=loaded_weight.narrow) as mock_narrow: + self.weight_loader(param, loaded_weight) + + # Verify narrow was called correctly: narrow(dim=0, start=0, length=2) + mock_narrow.assert_called_once_with(0, 0, 2) + + # Verify data was copied + self.assertTrue(torch.all(param == 1)) + + def test_weight_loader_tp_sharding_middle_rank(self): + """Test weight_loader with tensor parallelism sharding for middle rank""" + # Configure mocks for rank 2 of 4 + self.mock_tp_rank.return_value = 2 + self.mock_tp_size.return_value = 4 + + param = torch.zeros(2, 5) + loaded_weight = torch.ones(8, 5) + + with patch.object(loaded_weight, 'narrow', wraps=loaded_weight.narrow) as mock_narrow: + self.weight_loader(param, loaded_weight) + + # Verify narrow was called correctly: start = shard_size * rank = 2 * 2 = 4 + mock_narrow.assert_called_once_with(0, 4, 2) + + self.assertTrue(torch.all(param == 1)) + + def test_weight_loader_tp_sharding_last_rank(self): + """Test weight_loader with tensor parallelism sharding for last rank""" + # Configure mocks for rank 3 of 4 + self.mock_tp_rank.return_value = 3 + self.mock_tp_size.return_value = 4 + + param = torch.zeros(2, 5) + loaded_weight = torch.ones(8, 5) + + with patch.object(loaded_weight, 'narrow', wraps=loaded_weight.narrow) as mock_narrow: + self.weight_loader(param, loaded_weight) + + # Verify narrow was called correctly: start = 2 * 3 = 6 + mock_narrow.assert_called_once_with(0, 6, 2) + + self.assertTrue(torch.all(param == 1)) + + def test_weight_loader_shape_mismatch(self): + """Test weight_loader raises assertion error on shape mismatch""" + self.mock_tp_rank.return_value = 0 + self.mock_tp_size.return_value = 2 + + param = torch.zeros(2, 3) + loaded_weight = torch.ones(4, 4) # Different shape after sharding + + # Mock narrow to return tensor with wrong shape + with patch.object(loaded_weight, 'narrow', return_value=torch.ones(2, 4)): + with self.assertRaises(AssertionError) as context: + self.weight_loader(param, loaded_weight) + + # Verify error message contains expected information + self.assertIn("Attempted to load weight", str(context.exception)) + self.assertIn("into parameter", str(context.exception)) + + def test_weight_loader_with_different_dtypes(self): + """Test weight_loader handles different dtypes correctly""" + self.mock_tp_rank.return_value = 0 + self.mock_tp_size.return_value = 1 # No sharding + + param = torch.zeros(2, 3, dtype=torch.float32) + loaded_weight = torch.ones(2, 3, dtype=torch.float16) + + self.weight_loader(param, loaded_weight) + + # Verify data was copied and converted + self.assertTrue(torch.all(param == 1)) + self.assertEqual(param.dtype, torch.float32) + + +class TestAscendFAQuantAttentionMethodInit(unittest.TestCase): + """Test cases for AscendFAQuantAttentionMethod initialization""" + + def setUp(self): + """Set up test environment""" + # Mock vllm_config + self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config") + self.mock_get_config = self.config_patch.start() + + # Create mock config with attributes + self.mock_config = Mock() + self.mock_hf_config = Mock() + self.mock_hf_config.kv_lora_rank = 128 + self.mock_hf_config.qk_rope_head_dim = 64 + self.mock_config.model_config.hf_config = self.mock_hf_config + self.mock_get_config.return_value = self.mock_config + + # Import the class after patching + from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod + self.method_class = AscendFAQuantAttentionMethod + + def tearDown(self): + """Clean up after each test""" + self.config_patch.stop() + + def test_init_with_full_config(self): + """Test initialization when config has all attributes""" + method = self.method_class() + + self.assertTrue(method.transpose_weight) + self.assertFalse(method.printFlag) + self.assertEqual(method.kv_lora_rank, 128) + self.assertEqual(method.qk_rope_head_dim, 64) + + def test_init_without_kv_lora_rank(self): + """Test initialization when config lacks kv_lora_rank""" + delattr(self.mock_hf_config, "kv_lora_rank") + + method = self.method_class() + + self.assertEqual(method.kv_lora_rank, 0) + self.assertEqual(method.qk_rope_head_dim, 64) + + def test_init_without_qk_rope_head_dim(self): + """Test initialization when config lacks qk_rope_head_dim""" + delattr(self.mock_hf_config, "qk_rope_head_dim") + + method = self.method_class() + + self.assertEqual(method.kv_lora_rank, 128) + self.assertEqual(method.qk_rope_head_dim, 0) + + def test_init_without_both_attributes(self): + """Test initialization when config lacks both attributes""" + delattr(self.mock_hf_config, "kv_lora_rank") + delattr(self.mock_hf_config, "qk_rope_head_dim") + + method = self.method_class() + + self.assertEqual(method.kv_lora_rank, 0) + self.assertEqual(method.qk_rope_head_dim, 0) + + +class TestAscendFAQuantAttentionMethodCreateWeights(unittest.TestCase): + """Test cases for create_weights method""" + + def setUp(self): + """Set up test environment""" + # Mock vllm_config + self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config") + self.mock_get_config = self.config_patch.start() + + self.mock_config = Mock() + self.mock_hf_config = Mock() + self.mock_hf_config.kv_lora_rank = 128 + self.mock_hf_config.qk_rope_head_dim = 64 + self.mock_config.model_config.hf_config = self.mock_hf_config + self.mock_get_config.return_value = self.mock_config + + # Import the class + from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod + self.method_class = AscendFAQuantAttentionMethod + + # Mock torch functions + self.default_dtype_patch = patch("torch.get_default_dtype", return_value=torch.float32) + self.mock_default_dtype = self.default_dtype_patch.start() + + # Create a real nn.Module for testing + self.layer = nn.Module() + self.layer.num_heads = 32 + self.layer.num_kv_heads = 1 + + def tearDown(self): + """Clean up after each test""" + self.config_patch.stop() + self.default_dtype_patch.stop() + + def test_create_weights_adds_submodules(self): + """Test that create_weights adds fa_q, fa_k, fa_v submodules""" + method = self.method_class() + + with patch("torch.empty") as mock_empty: + mock_empty.return_value = torch.zeros(1, 1) + + method.create_weights(self.layer) + + # Verify submodules were added + self.assertTrue(hasattr(self.layer, "fa_q")) + self.assertTrue(hasattr(self.layer, "fa_k")) + self.assertTrue(hasattr(self.layer, "fa_v")) + + # Verify they are instances of nn.Module + self.assertIsInstance(self.layer.fa_q, nn.Module) + self.assertIsInstance(self.layer.fa_k, nn.Module) + self.assertIsInstance(self.layer.fa_v, nn.Module) + + def test_create_weights_creates_correct_tensors(self): + """Test that create_weights creates tensors with correct shapes and dtypes""" + method = self.method_class() + + # Track torch.empty calls + empty_calls = [] + + def mock_empty(size, dtype=None): + empty_calls.append((size, dtype)) + return torch.zeros(size, dtype=dtype if dtype else torch.float32) + + with patch("torch.empty", side_effect=mock_empty): + method.create_weights(self.layer) + + # Verify tensor creations + expected_calls = [ + ((32, 1), torch.float32), # fa_q.scale + ((1, 1), torch.float32), # fa_k.scale + ((1, 1), torch.float32), # fa_v.scale + ((32, 1), torch.int8), # fa_q.offset + ((1, 1), torch.int8), # fa_k.offset + ((1, 1), torch.int8), # fa_v.offset + ] + + # Compare without considering order + self.assertEqual(len(empty_calls), len(expected_calls)) + for call in expected_calls: + self.assertIn(call, empty_calls) + + def test_create_weights_registers_parameters(self): + """Test that create_weights registers parameters with correct attributes""" + method = self.method_class() + + # Create real tensors for testing + def create_tensor(*args, **kwargs): + size = args[0] if args else kwargs.get('size', (1,)) + dtype = kwargs.get('dtype', torch.float32) + return torch.zeros(*size, dtype=dtype) + + with patch("torch.empty", side_effect=create_tensor): + method.create_weights(self.layer) + + # Import weight_loader for comparison + from vllm_ascend.quantization.methods.kv_c8 import _fa_quant_weight_loader as weight_loader + + # Verify each parameter exists and has weight_loader + self.assertTrue(hasattr(self.layer.fa_q, "scale")) + self.assertTrue(hasattr(self.layer.fa_q.scale, "weight_loader")) + self.assertEqual(self.layer.fa_q.scale.weight_loader, weight_loader) + self.assertFalse(self.layer.fa_q.scale.requires_grad) + + self.assertTrue(hasattr(self.layer.fa_k, "scale")) + self.assertTrue(hasattr(self.layer.fa_k.scale, "weight_loader")) + + self.assertTrue(hasattr(self.layer.fa_v, "scale")) + self.assertTrue(hasattr(self.layer.fa_v.scale, "weight_loader")) + + self.assertTrue(hasattr(self.layer.fa_q, "offset")) + self.assertTrue(hasattr(self.layer.fa_q.offset, "weight_loader")) + self.assertEqual(self.layer.fa_q.offset.dtype, torch.int8) + + +class TestAscendFAQuantAttentionMethodProcessWeights(unittest.TestCase): + """Test cases for process_weights_after_loading method""" + + def setUp(self): + """Set up test environment""" + # Mock vllm_config + self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config") + self.mock_get_config = self.config_patch.start() + + self.mock_config = Mock() + self.mock_hf_config = Mock() + self.mock_hf_config.kv_lora_rank = 64 + self.mock_hf_config.qk_rope_head_dim = 32 + self.mock_config.model_config.hf_config = self.mock_hf_config + self.mock_get_config.return_value = self.mock_config + + # Import the class + from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod + self.method_class = AscendFAQuantAttentionMethod + + # Create method instance with real layer + self.method = self.method_class() + + # Create a real nn.Module for testing + self.layer = nn.Module() + + # Create real tensors for fa_k + self.fa_k_scale = torch.tensor([[2.0, 3.0, 4.0]], dtype=torch.float16) # Shape (1,3) + self.fa_k_offset = torch.tensor([[1, 2, 3]], dtype=torch.int8) # Shape (1,3) + + # Create fa_k module with parameters + self.layer.fa_k = nn.Module() + self.layer.fa_k.scale = nn.Parameter(self.fa_k_scale, requires_grad=False) + self.layer.fa_k.offset = nn.Parameter(self.fa_k_offset, requires_grad=False) + + def tearDown(self): + """Clean up after each test""" + self.config_patch.stop() + + def test_process_weights_with_single_value_scale(self): + """Test process_weights with single value scale""" + # Create new layer with single value scale + layer = nn.Module() + layer.fa_k = nn.Module() + layer.fa_k.scale = nn.Parameter(torch.tensor([[2.0]], dtype=torch.float16), requires_grad=False) + layer.fa_k.offset = nn.Parameter(torch.tensor([[1]], dtype=torch.int8), requires_grad=False) + + self.method.kv_lora_rank = 4 + self.method.process_weights_after_loading(layer) + + self.assertEqual(layer.quant_kscale.shape, (1, 4)) + self.assertEqual(layer.quant_kscale.dtype, torch.float32) + + +class TestIntegration(unittest.TestCase): + """Integration tests for the complete kv_c8 functionality""" + + def setUp(self): + """Set up test environment""" + # Mock vllm_config + self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config") + self.mock_get_config = self.config_patch.start() + + self.mock_config = Mock() + self.mock_hf_config = Mock() + self.mock_hf_config.kv_lora_rank = 64 + self.mock_hf_config.qk_rope_head_dim = 32 + self.mock_config.model_config.hf_config = self.mock_hf_config + self.mock_get_config.return_value = self.mock_config + + # Mock distributed functions + self.tp_rank_patch = patch( + "vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_rank" + ) + self.tp_size_patch = patch( + "vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_world_size" + ) + self.mock_tp_rank = self.tp_rank_patch.start() + self.mock_tp_size = self.tp_size_patch.start() + + def tearDown(self): + """Clean up after each test""" + self.config_patch.stop() + self.tp_rank_patch.stop() + self.tp_size_patch.stop() + + def test_complete_workflow(self): + """Test complete workflow from weight creation to processing""" + from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod + + # Create method instance + method = AscendFAQuantAttentionMethod() + + # Create real layer + layer = nn.Module() + layer.num_heads = 32 + layer.num_kv_heads = 1 + + # Step 1: Create weights + method.create_weights(layer) + + # Verify weights were created with correct structure + self.assertTrue(hasattr(layer, "fa_q")) + self.assertTrue(hasattr(layer, "fa_k")) + self.assertTrue(hasattr(layer, "fa_v")) + + self.assertTrue(hasattr(layer.fa_q, "scale")) + self.assertTrue(hasattr(layer.fa_q, "offset")) + self.assertTrue(hasattr(layer.fa_k, "scale")) + self.assertTrue(hasattr(layer.fa_k, "offset")) + self.assertTrue(hasattr(layer.fa_v, "scale")) + self.assertTrue(hasattr(layer.fa_v, "offset")) + + # Step 2: Simulate weight loading + self.mock_tp_rank.return_value = 0 + self.mock_tp_size.return_value = 1 + + # Create dummy weights + q_scale = torch.randn(32, 1) + k_scale = torch.randn(1, 1) + v_scale = torch.randn(1, 1) + q_offset = torch.randint(-128, 127, (32, 1), dtype=torch.int8) + k_offset = torch.randint(-128, 127, (1, 1), dtype=torch.int8) + v_offset = torch.randint(-128, 127, (1, 1), dtype=torch.int8) + + # Load weights using weight_loader + from vllm_ascend.quantization.methods.kv_c8 import _fa_quant_weight_loader as weight_loader + + with torch.no_grad(): + weight_loader(layer.fa_q.scale, q_scale) + weight_loader(layer.fa_k.scale, k_scale) + weight_loader(layer.fa_v.scale, v_scale) + weight_loader(layer.fa_q.offset, q_offset) + weight_loader(layer.fa_k.offset, k_offset) + weight_loader(layer.fa_v.offset, v_offset) + + # Verify weights were loaded correctly + self.assertTrue(torch.all(layer.fa_q.scale == q_scale)) + self.assertTrue(torch.all(layer.fa_k.scale == k_scale)) + self.assertTrue(torch.all(layer.fa_v.scale == v_scale)) + + # Step 3: Process after loading + method.process_weights_after_loading(layer) + + # Verify processed parameters + self.assertTrue(hasattr(layer, "fak_descale")) + self.assertTrue(hasattr(layer, "fak_offset")) + self.assertTrue(hasattr(layer, "quant_kscale")) + + +class TestC8KVScaleWeightLoader(TestBase): + """Tests for _c8_kv_scale_weight_loader in kv_c8.py.""" + + def setUp(self): + from vllm_ascend.quantization.methods.kv_c8 import _c8_kv_scale_weight_loader + self.loader = _c8_kv_scale_weight_loader + + def test_shape_match_copies_value(self): + param = nn.Parameter(torch.ones(4, dtype=torch.float32), requires_grad=False) + loaded = torch.tensor([1.0, 2.0, 3.0, 4.0]) + self.loader(param, loaded) + self.assertTrue(torch.allclose(param.data, loaded.float())) + + def test_shape_mismatch_resizes_param(self): + param = nn.Parameter(torch.ones(1, dtype=torch.float32), requires_grad=False) + loaded = torch.arange(8, dtype=torch.float32) + self.loader(param, loaded) + self.assertEqual(param.data.shape, (8,)) + self.assertTrue(torch.allclose(param.data, loaded)) + + def test_squeeze_before_compare(self): + param = nn.Parameter(torch.ones(4, dtype=torch.float32), requires_grad=False) + loaded = torch.arange(4, dtype=torch.float32).unsqueeze(0) # shape [1, 4] + self.loader(param, loaded) + self.assertEqual(param.data.shape, (4,)) + + def test_dtype_preserved_as_param_dtype(self): + param = nn.Parameter(torch.ones(4, dtype=torch.float32), requires_grad=False) + loaded = torch.arange(4, dtype=torch.float16) + self.loader(param, loaded) + self.assertEqual(param.data.dtype, torch.float32) + + +class TestAscendC8KVCacheAttentionMethod(TestBase): + """Tests for AscendC8KVCacheAttentionMethod in kv_c8.py.""" + + def _make_method(self): + from vllm_ascend.quantization.methods.kv_c8 import AscendC8KVCacheAttentionMethod + return AscendC8KVCacheAttentionMethod(quant_description={}, prefix="model.layers.0.self_attn.attn") + + def _make_layer_with_impl(self): + layer = nn.Module() + layer.impl = MagicMock() + return layer + + def test_create_weights_sets_kv_cache_torch_dtype(self): + method = self._make_method() + layer = self._make_layer_with_impl() + method.create_weights(layer) + self.assertEqual(layer.kv_cache_torch_dtype, torch.int8) + + def test_create_weights_registers_scale_offset_params(self): + method = self._make_method() + layer = self._make_layer_with_impl() + method.create_weights(layer) + self.assertIsInstance(layer.k_cache_scale, nn.Parameter) + self.assertIsInstance(layer.k_cache_offset, nn.Parameter) + self.assertIsInstance(layer.v_cache_scale, nn.Parameter) + self.assertIsInstance(layer.v_cache_offset, nn.Parameter) + self.assertFalse(layer.k_cache_scale.requires_grad) + self.assertFalse(layer.v_cache_offset.requires_grad) + + def test_create_weights_initial_values(self): + method = self._make_method() + layer = self._make_layer_with_impl() + method.create_weights(layer) + self.assertEqual(layer.k_cache_scale.data.item(), 1.0) + self.assertEqual(layer.v_cache_scale.data.item(), 1.0) + self.assertEqual(layer.k_cache_offset.data.item(), 0.0) + self.assertEqual(layer.v_cache_offset.data.item(), 0.0) + + def test_create_weights_assigns_weight_loader(self): + from vllm_ascend.quantization.methods.kv_c8 import _c8_kv_scale_weight_loader + method = self._make_method() + layer = self._make_layer_with_impl() + method.create_weights(layer) + self.assertIs(layer.k_cache_scale.weight_loader, _c8_kv_scale_weight_loader) + self.assertIs(layer.v_cache_scale.weight_loader, _c8_kv_scale_weight_loader) + self.assertIs(layer.k_cache_offset.weight_loader, _c8_kv_scale_weight_loader) + self.assertIs(layer.v_cache_offset.weight_loader, _c8_kv_scale_weight_loader) + + def test_process_weights_after_loading_flattens(self): + method = self._make_method() + layer = nn.Module() + layer.k_cache_scale = nn.Parameter(torch.ones(2, 4), requires_grad=False) + layer.k_cache_offset = nn.Parameter(torch.zeros(2, 4), requires_grad=False) + layer.v_cache_scale = nn.Parameter(torch.ones(2, 4), requires_grad=False) + layer.v_cache_offset = nn.Parameter(torch.zeros(2, 4), requires_grad=False) + method.process_weights_after_loading(layer) + self.assertEqual(layer.k_cache_scale.data.dim(), 1) + self.assertEqual(layer.k_cache_scale.data.shape[0], 8) + self.assertEqual(layer.v_cache_offset.data.dim(), 1) + + def test_apply_raises_runtime_error(self): + method = self._make_method() + layer = MagicMock() + with self.assertRaises(RuntimeError): + method.apply(layer, MagicMock(), MagicMock(), MagicMock(), None, None, None, None, None) + + +class TestAscendC8AttentionBackendImplScales(TestBase): + """Tests for AscendC8AttentionBackendImpl scale helpers.""" + + def _make_impl(self, num_kv_heads=4, head_size=8): + from vllm_ascend.attention.attention_v1 import AscendC8AttentionBackendImpl + impl = object.__new__(AscendC8AttentionBackendImpl) + impl.num_heads = num_kv_heads + impl.num_kv_heads = num_kv_heads + impl.head_size = head_size + impl.scale = 1.0 + impl.key_cache = None + impl.value_cache = None + return impl + + def _make_layer(self, num_kv_heads=4, head_size=8): + layer = nn.Module() + layer.k_cache_scale = nn.Parameter( + torch.ones(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False + ) + layer.k_cache_offset = nn.Parameter( + torch.zeros(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False + ) + layer.v_cache_scale = nn.Parameter( + torch.ones(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False + ) + layer.v_cache_offset = nn.Parameter( + torch.zeros(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False + ) + return layer + + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) + def test_prepare_c8_scales_runs_once(self, mock_tp_size, mock_tp_rank): + impl = self._make_impl() + layer = self._make_layer() + impl._prepare_c8_scales(layer, torch.device("cpu")) + self.assertTrue(hasattr(layer, "_c8_scales_prepared")) + self.assertTrue(layer._c8_scales_prepared) + + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) + def test_prepare_c8_scales_idempotent(self, mock_tp_size, mock_tp_rank): + impl = self._make_impl() + layer = self._make_layer() + impl._prepare_c8_scales(layer, torch.device("cpu")) + k_scale_after_first = layer._c8_k_scale.clone() + layer.k_cache_scale.data = torch.ones(32, dtype=torch.float32) * 99 + impl._prepare_c8_scales(layer, torch.device("cpu")) + self.assertTrue(torch.allclose(layer._c8_k_scale, k_scale_after_first)) + + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) + def test_prepare_c8_scales_creates_bnsd_shape(self, mock_tp_size, mock_tp_rank): + num_kv_heads, head_size = 4, 8 + impl = self._make_impl(num_kv_heads, head_size) + layer = self._make_layer(num_kv_heads, head_size) + impl._prepare_c8_scales(layer, torch.device("cpu")) + self.assertEqual(layer._c8_k_aq_scale.shape, (1, num_kv_heads, 1, head_size)) + self.assertEqual(layer._c8_v_aq_scale.shape, (1, num_kv_heads, 1, head_size)) + self.assertEqual(layer._c8_k_aq_scale.dtype, torch.bfloat16) + + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) + def test_quantize_kv_to_int8_output_dtype(self, mock_tp_size, mock_tp_rank): + num_kv_heads, head_size = 4, 8 + impl = self._make_impl(num_kv_heads, head_size) + layer = self._make_layer(num_kv_heads, head_size) + impl._prepare_c8_scales(layer, torch.device("cpu")) + num_tokens = 6 + key = torch.zeros(num_tokens, num_kv_heads, head_size, dtype=torch.bfloat16) + value = torch.zeros(num_tokens, num_kv_heads, head_size, dtype=torch.bfloat16) + key_q, value_q = impl._quantize_kv_to_int8(key, value, layer, num_tokens) + self.assertEqual(key_q.dtype, torch.int8) + self.assertEqual(value_q.dtype, torch.int8) + self.assertEqual(key_q.shape, key.shape) + + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) + def test_quantize_kv_to_int8_formula(self, mock_tp_size, mock_tp_rank): + """With scale=2.0, offset=0: q = round(x / 2).""" + num_kv_heads, head_size = 1, 4 + impl = self._make_impl(num_kv_heads, head_size) + layer = nn.Module() + scale_val = torch.full((num_kv_heads * head_size,), 2.0, dtype=torch.float32) + layer.k_cache_scale = nn.Parameter(scale_val.clone(), requires_grad=False) + layer.k_cache_offset = nn.Parameter(torch.zeros(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False) + layer.v_cache_scale = nn.Parameter(scale_val.clone(), requires_grad=False) + layer.v_cache_offset = nn.Parameter(torch.zeros(num_kv_heads * head_size, dtype=torch.float32), requires_grad=False) + impl._prepare_c8_scales(layer, torch.device("cpu")) + key = torch.full((1, num_kv_heads, head_size), 4.0, dtype=torch.bfloat16) + value = torch.full((1, num_kv_heads, head_size), 4.0, dtype=torch.bfloat16) + key_q, _ = impl._quantize_kv_to_int8(key, value, layer, 1) + self.assertTrue(torch.all(key_q[0] == 2)) + + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0) + @patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1) + def test_dequant_paged_kv_to_dense_round_trip(self, mock_tp_size, mock_tp_rank): + """With scale=1, offset=0: dequant(int8) == float(int8).""" + num_kv_heads, head_size = 2, 4 + block_size = 32 + num_blocks = 2 + H = num_kv_heads * head_size + impl = self._make_impl(num_kv_heads, head_size) + layer = self._make_layer(num_kv_heads, head_size) + impl._prepare_c8_scales(layer, torch.device("cpu")) + + key_int8 = torch.randint(-10, 10, (num_blocks, block_size, H), dtype=torch.int8) + value_int8 = torch.randint(-10, 10, (num_blocks, block_size, H), dtype=torch.int8) + seq_lens = [32, 32] + block_table = torch.tensor([[0], [1]], dtype=torch.long) + + dense_k, dense_v = impl._dequant_paged_kv_to_dense( + key_int8, value_int8, block_table, seq_lens, torch.float32, layer + ) + expected_k = key_int8.view(-1, num_kv_heads, head_size).float() + self.assertEqual(dense_k.shape, (64, num_kv_heads, head_size)) + self.assertTrue(torch.allclose(dense_k, expected_k)) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/ut/quantization/methods/test_registry.py b/tests/ut/quantization/methods/test_registry.py new file mode 100644 index 000000000000..ecb99a05077c --- /dev/null +++ b/tests/ut/quantization/methods/test_registry.py @@ -0,0 +1,83 @@ +from tests.ut.base import TestBase +from vllm_ascend.quantization.methods.base import ( + AscendLinearScheme, + AscendMoEScheme, +) +from vllm_ascend.quantization.methods.registry import ( + _SCHEME_REGISTRY, + register_scheme, + get_scheme_class, +) + + +class TestRegisterScheme(TestBase): + + def test_register_scheme(self): + @register_scheme("TEST_QUANT_TYPE", "linear") + class TestLinearScheme(AscendLinearScheme): + def get_weight(self, input_size, output_size, params_dtype): + return {} + + def apply(self, layer, x, bias=None, tp_rank=0): + return x + + scheme_class = get_scheme_class("TEST_QUANT_TYPE", "linear") + self.assertIs(scheme_class, TestLinearScheme) + + +class TestGetSchemeClass(TestBase): + + def test_get_scheme_class_existing_linear(self): + cls = get_scheme_class("W8A8_DYNAMIC", "linear") + self.assertIsNotNone(cls) + self.assertTrue(issubclass(cls, AscendLinearScheme)) + + def test_get_scheme_class_existing_moe(self): + cls = get_scheme_class("W8A8_DYNAMIC", "moe") + self.assertIsNotNone(cls) + self.assertTrue(issubclass(cls, AscendMoEScheme)) + + def test_get_scheme_class_existing_attention(self): + cls = get_scheme_class("FAKQuant", "attention") + self.assertIsNotNone(cls) + + def test_get_scheme_class_nonexistent(self): + cls = get_scheme_class("NONEXISTENT", "linear") + self.assertIsNone(cls) + + def test_get_scheme_class_nonexistent_layer_type(self): + cls = get_scheme_class("W8A8_DYNAMIC", "nonexistent") + self.assertIsNone(cls) + + def test_register_scheme_duplicate_raises(self): + with self.assertRaises(ValueError): + @register_scheme("W8A8_DYNAMIC", "linear") + class Duplicate: + pass + + def test_all_linear_schemes_subclass_ascend_linear_scheme(self): + for (quant_type, layer_type), scheme_cls in _SCHEME_REGISTRY.items(): + if layer_type == "linear": + self.assertTrue( + issubclass(scheme_cls, AscendLinearScheme), + f"{scheme_cls.__name__} for {quant_type}/{layer_type} " + f"should be subclass of AscendLinearScheme", + ) + + def test_all_moe_schemes_subclass_ascend_moe_scheme(self): + for (quant_type, layer_type), scheme_cls in _SCHEME_REGISTRY.items(): + if layer_type == "moe": + self.assertTrue( + issubclass(scheme_cls, AscendMoEScheme), + f"{scheme_cls.__name__} for {quant_type}/{layer_type} " + f"should be subclass of AscendMoEScheme", + ) + + def test_registry_not_empty(self): + self.assertGreater(len(_SCHEME_REGISTRY), 0) + + def test_registry_key_format(self): + for key in _SCHEME_REGISTRY.keys(): + self.assertEqual(len(key), 2) + self.assertIsInstance(key[0], str) + self.assertIsInstance(key[1], str) \ No newline at end of file diff --git a/tests/ut/quantization/test_w4a16.py b/tests/ut/quantization/methods/test_w4a16.py similarity index 95% rename from tests/ut/quantization/test_w4a16.py rename to tests/ut/quantization/methods/test_w4a16.py index 87c2b79ffe00..b4df43fd2a99 100644 --- a/tests/ut/quantization/test_w4a16.py +++ b/tests/ut/quantization/methods/test_w4a16.py @@ -1,307 +1,317 @@ -from unittest.mock import Mock, patch - -import torch - -from tests.ut.base import TestBase -from vllm_ascend.ascend_forward_context import MoECommType -from vllm_ascend.quantization.methods.w4a16 import AscendW4A16FusedMoEMethod, pack_to_int32, unpack_from_int32 - - -class TestUnpackFromInt32(TestBase): - - def test_unpack_from_int32_packed_dim_1(self): - weight = torch.tensor([[305419896, -1420531520]], dtype=torch.int32) - shape = torch.Size([1, 8]) - num_bits = 4 - - result = unpack_from_int32(weight, shape, num_bits, packed_dim=1) - - self.assertEqual(result.dtype, torch.int8) - self.assertEqual(result.shape, shape) - - def test_unpack_from_int32_packed_dim_0(self): - weight = torch.tensor([[305419896], [-1420531520]], dtype=torch.int32) - shape = torch.Size([8, 1]) - num_bits = 4 - - result = unpack_from_int32(weight, shape, num_bits, packed_dim=0) - - self.assertEqual(result.dtype, torch.int8) - self.assertEqual(result.shape, shape) - - def test_unpack_from_int32_assertions(self): - with self.assertRaises(AssertionError): - weight = torch.tensor([[1, 2]], dtype=torch.int64) - unpack_from_int32(weight, torch.Size([8, 1]), 4) - - with self.assertRaises(AssertionError): - weight = torch.tensor([[1, 2]], dtype=torch.int32) - unpack_from_int32(weight, torch.Size([8, 1]), 16) - - -class TestPackToInt32(TestBase): - - @patch( - "vllm_ascend.quantization.methods.w4a16.torch_npu.npu_convert_weight_to_int4pack" - ) - def test_pack_to_int32_int8(self, mock_npu_convert_weight_to_int4pack): - mock_npu_convert_weight_to_int4pack.return_value = torch.zeros( - (2, 4), dtype=torch.int32) - - weight = torch.zeros((2, 8, 16), dtype=torch.int8) - result = pack_to_int32(weight) - - self.assertEqual(result.dtype, torch.int32) - mock_npu_convert_weight_to_int4pack.assert_not_called() - - self.assertEqual(result.shape, torch.Size([2, 8, 4])) - - @patch( - "vllm_ascend.quantization.methods.w4a16.torch_npu.npu_convert_weight_to_int4pack" - ) - def test_pack_to_int32_int32(self, mock_npu_convert_weight_to_int4pack): - - def mock_convert_weight(weight): - return weight - - mock_npu_convert_weight_to_int4pack.side_effect = mock_convert_weight - weight = torch.zeros((2, 8, 8), dtype=torch.int32) - result = pack_to_int32(weight) - - self.assertEqual(result.dtype, torch.int32) - self.assertEqual(result.shape, weight.shape) - - def test_pack_to_int32_assertion_dim(self): - with self.assertRaises(AssertionError): - weight = torch.zeros((8, 8), dtype=torch.int8) - pack_to_int32(weight) - - def test_pack_to_int32_assertion_dtype(self): - with self.assertRaises(AssertionError): - weight = torch.zeros((2, 8, 8), dtype=torch.float32) - pack_to_int32(weight) - - def test_pack_to_int32_assertion_divisible(self): - with self.assertRaises(AssertionError): - weight = torch.zeros((2, 8, 7), dtype=torch.int32) - pack_to_int32(weight) - - with self.assertRaises(AssertionError): - weight = torch.zeros((2, 8, 7), dtype=torch.int8) - pack_to_int32(weight) - - -class TestAscendW4A16FusedMoEMethod(TestBase): - experts = 8 - input_size = 32 - output_size = 128 - group_size = 32 - - @patch("vllm_ascend.quantization.methods.w4a16.get_ascend_config") - @patch("vllm_ascend.quantization.methods.w4a16.get_current_vllm_config") - def setUp(self, mock_get_current_vllm_config, mock_get_ascend_config): - mock_ascend_config = Mock() - mock_ascend_config.eplb_config.dynamic_eplb = False - mock_ascend_config.eplb_config.expert_map_record_path = None - 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, - }) - mock_get_current_vllm_config.return_value = mock_vllm_config - - self.quant_method = AscendW4A16FusedMoEMethod() - - def test_init(self): - self.assertTrue(self.quant_method.transpose_weight) - self.assertEqual(self.quant_method.num_bits, 4) - self.assertEqual(self.quant_method.pack_factor, 8) - self.assertEqual(self.quant_method.group_size, self.group_size) - self.assertFalse(self.quant_method.dynamic_eplb) - - def test_get_weight(self): - param_dict = self.quant_method.get_weight(self.experts, - self.input_size, - self.output_size, - torch.bfloat16) - - self.assertEqual(param_dict["w13_weight_packed"].dtype, torch.int32) - expected_w13_shape = (self.experts, 2 * self.input_size, - self.output_size // - self.quant_method.pack_factor) - self.assertEqual(param_dict["w13_weight_packed"].shape, - expected_w13_shape) - - self.assertEqual(param_dict["w2_weight_packed"].dtype, torch.int32) - expected_w2_shape = (self.experts, self.output_size, - self.input_size // self.quant_method.pack_factor) - self.assertEqual(param_dict["w2_weight_packed"].shape, - expected_w2_shape) - - def test_get_dynamic_quant_param(self): - 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.bfloat16) - expected_w13_scale_shape = (self.experts, 2 * self.input_size, - self.output_size // self.group_size) - self.assertEqual(param_dict["w13_weight_scale"].shape, - expected_w13_scale_shape) - - self.assertEqual(param_dict["w2_weight_scale"].dtype, torch.bfloat16) - expected_w2_scale_shape = (self.experts, self.output_size, - self.input_size // self.group_size) - self.assertEqual(param_dict["w2_weight_scale"].shape, - expected_w2_scale_shape) - - self.assertEqual(param_dict["w13_weight_shape"].dtype, torch.int32) - self.assertEqual(param_dict["w13_weight_shape"].shape, - (self.experts, 2)) - - self.assertEqual(param_dict["w2_weight_shape"].dtype, torch.int32) - self.assertEqual(param_dict["w2_weight_shape"].shape, - (self.experts, 2)) - - self.assertEqual(param_dict["w13_weight_offset"].dtype, torch.bfloat16) - self.assertEqual(param_dict["w13_weight_offset"].shape, - expected_w13_scale_shape) - - self.assertEqual(param_dict["w2_weight_offset"].dtype, torch.bfloat16) - self.assertEqual(param_dict["w2_weight_offset"].shape, - expected_w2_scale_shape) - - def build_layer(self): - """Build a mock layer for testing""" - layer = torch.nn.Module() - - w13_shape = (self.experts, 2 * self.input_size, - self.output_size // self.quant_method.pack_factor) - w2_shape = (self.experts, self.output_size, - self.input_size // self.quant_method.pack_factor) - - layer.w13_weight_packed = torch.nn.Parameter(torch.randint( - -100, 100, w13_shape, dtype=torch.int32), - requires_grad=False) - layer.w2_weight_packed = torch.nn.Parameter(torch.randint( - -100, 100, w2_shape, dtype=torch.int32), - requires_grad=False) - - w13_scale_shape = (self.experts, 2 * self.input_size, - self.output_size // self.group_size) - w2_scale_shape = (self.experts, self.output_size, - self.input_size // self.group_size) - - layer.w13_weight_scale = torch.nn.Parameter(torch.ones( - w13_scale_shape, dtype=torch.bfloat16), - requires_grad=False) - layer.w2_weight_scale = torch.nn.Parameter(torch.ones( - w2_scale_shape, dtype=torch.bfloat16), - requires_grad=False) - - layer.w13_weight_offset = torch.nn.Parameter(torch.zeros( - w13_scale_shape, dtype=torch.bfloat16), - requires_grad=False) - layer.w2_weight_offset = torch.nn.Parameter(torch.zeros( - w2_scale_shape, dtype=torch.bfloat16), - requires_grad=False) - - layer.w13_weight_shape = torch.nn.Parameter(torch.tensor( - [[2 * self.input_size, self.output_size]] * self.experts, - dtype=torch.int32), - requires_grad=False) - layer.w2_weight_shape = torch.nn.Parameter(torch.tensor( - [[self.output_size, self.input_size]] * self.experts, - dtype=torch.int32), - requires_grad=False) - - return layer - - @patch( - "vllm_ascend.quantization.methods.w4a16.torch_npu.npu_convert_weight_to_int4pack" - ) - def test_process_weights_after_loading_with_transpose( - self, mock_npu_convert_weight_to_int4pack): - - def mock_convert_weight(weight): - new_shape = list(weight.shape) - new_shape[-1] = new_shape[-1] // 8 - return torch.zeros(new_shape, dtype=torch.int32) - - mock_npu_convert_weight_to_int4pack.side_effect = mock_convert_weight - - layer = self.build_layer() - self.quant_method.transpose_weight = True - - self.quant_method.process_weights_after_loading(layer) - - self.assertEqual(layer.w13_weight_packed.data.shape, - torch.Size([8, 128, 8])) - self.assertEqual(layer.w2_weight_packed.data.shape, - torch.Size([8, 32, 16])) - - self.assertEqual(layer.w13_weight_scale.data.shape, - torch.Size([8, 4, 64])) - self.assertEqual(layer.w2_weight_scale.data.shape, - torch.Size([8, 1, 128])) - self.assertEqual(layer.w13_weight_offset.data.shape, - torch.Size([8, 4, 64])) - self.assertEqual(layer.w2_weight_offset.data.shape, - torch.Size([8, 1, 128])) - - self.assertTrue(layer.w13_weight_scale.data.is_contiguous()) - self.assertTrue(layer.w2_weight_scale.data.is_contiguous()) - self.assertTrue(layer.w13_weight_offset.data.is_contiguous()) - self.assertTrue(layer.w2_weight_offset.data.is_contiguous()) - - def test_process_weights_after_loading_without_transpose(self): - layer = self.build_layer() - self.quant_method.transpose_weight = False - - original_w13_data = layer.w13_weight_packed.data.clone() - original_w2_data = layer.w2_weight_packed.data.clone() - - self.quant_method.process_weights_after_loading(layer) - - self.assertTrue( - torch.equal(layer.w13_weight_packed.data, original_w13_data)) - self.assertTrue( - torch.equal(layer.w2_weight_packed.data, original_w2_data)) - - @patch("vllm_ascend.quantization.methods.w4a16._EXTRA_CTX") - @patch("vllm_ascend.quantization.methods.w4a16.select_experts") - def test_apply_uses_explicit_dispatch_and_mlp_args(self, mock_select_experts, mock_extra_ctx): - tokens = 3 - hidden_size = self.output_size - layer = self.build_layer() - x = torch.randn(tokens, hidden_size, dtype=torch.float32) - router_logits = torch.randn(tokens, self.experts, dtype=torch.float32) - topk_weights = torch.randn(tokens, 2, dtype=torch.float32) - topk_ids = torch.randint(0, self.experts, (tokens, 2), dtype=torch.int64) - mc2_mask = torch.tensor([1, 0, 1], dtype=torch.bool) - pertoken_scale = torch.randn(tokens, dtype=torch.float32) - - mock_select_experts.return_value = (topk_weights, topk_ids) - mock_comm = Mock() - mock_comm.fused_experts.return_value = torch.randn(tokens, hidden_size, dtype=torch.float32) - mock_extra_ctx.moe_comm_method = mock_comm - mock_extra_ctx.moe_comm_type = MoECommType.ALLGATHER - - self.quant_method.apply( - layer=layer, - x=x, - router_logits=router_logits, - top_k=2, - renormalize=True, - global_num_experts=self.experts, - activation="gelu", - apply_router_weight_on_input=True, - mc2_mask=mc2_mask, - pertoken_scale=pertoken_scale, - ) - - fused_experts_input = mock_comm.fused_experts.call_args.kwargs["fused_experts_input"] - self.assertEqual(fused_experts_input.activation, "gelu") - self.assertTrue(fused_experts_input.routing.apply_router_weight_on_input) - self.assertIs(fused_experts_input.routing.mc2_mask, mc2_mask) - self.assertIs(fused_experts_input.routing.pertoken_scale, pertoken_scale) +from unittest.mock import Mock, patch + +import torch + +from tests.ut.base import TestBase +from vllm_ascend.ascend_forward_context import MoECommType +from vllm_ascend.quantization.methods.w4a16 import AscendW4A16FusedMoEMethod, pack_to_int32, unpack_from_int32 + + +class TestUnpackFromInt32(TestBase): + + def test_unpack_from_int32_packed_dim_1(self): + weight = torch.tensor([[305419896, -1420531520]], dtype=torch.int32) + shape = torch.Size([1, 8]) + num_bits = 4 + + result = unpack_from_int32(weight, shape, num_bits, packed_dim=1) + + self.assertEqual(result.dtype, torch.int8) + self.assertEqual(result.shape, shape) + + def test_unpack_from_int32_packed_dim_0(self): + weight = torch.tensor([[305419896], [-1420531520]], dtype=torch.int32) + shape = torch.Size([8, 1]) + num_bits = 4 + + result = unpack_from_int32(weight, shape, num_bits, packed_dim=0) + + self.assertEqual(result.dtype, torch.int8) + self.assertEqual(result.shape, shape) + + def test_unpack_from_int32_assertions(self): + with self.assertRaises(AssertionError): + weight = torch.tensor([[1, 2]], dtype=torch.int64) + unpack_from_int32(weight, torch.Size([8, 1]), 4) + + with self.assertRaises(AssertionError): + weight = torch.tensor([[1, 2]], dtype=torch.int32) + unpack_from_int32(weight, torch.Size([8, 1]), 16) + + +class TestPackToInt32(TestBase): + + @patch( + "vllm_ascend.quantization.methods.w4a16.torch_npu.npu_convert_weight_to_int4pack" + ) + def test_pack_to_int32_int8(self, mock_npu_convert_weight_to_int4pack): + mock_npu_convert_weight_to_int4pack.return_value = torch.zeros( + (2, 4), dtype=torch.int32) + + weight = torch.zeros((2, 8, 16), dtype=torch.int8) + result = pack_to_int32(weight) + + self.assertEqual(result.dtype, torch.int32) + mock_npu_convert_weight_to_int4pack.assert_not_called() + + self.assertEqual(result.shape, torch.Size([2, 8, 4])) + + @patch( + "vllm_ascend.quantization.methods.w4a16.torch_npu.npu_convert_weight_to_int4pack" + ) + def test_pack_to_int32_int32(self, mock_npu_convert_weight_to_int4pack): + + def mock_convert_weight(weight): + return weight + + mock_npu_convert_weight_to_int4pack.side_effect = mock_convert_weight + weight = torch.zeros((2, 8, 8), dtype=torch.int32) + result = pack_to_int32(weight) + + self.assertEqual(result.dtype, torch.int32) + self.assertEqual(result.shape, weight.shape) + + def test_pack_to_int32_assertion_dim(self): + with self.assertRaises(AssertionError): + weight = torch.zeros((8, 8), dtype=torch.int8) + pack_to_int32(weight) + + def test_pack_to_int32_assertion_dtype(self): + with self.assertRaises(AssertionError): + weight = torch.zeros((2, 8, 8), dtype=torch.float32) + pack_to_int32(weight) + + def test_pack_to_int32_assertion_divisible(self): + with self.assertRaises(AssertionError): + weight = torch.zeros((2, 8, 7), dtype=torch.int32) + pack_to_int32(weight) + + with self.assertRaises(AssertionError): + weight = torch.zeros((2, 8, 7), dtype=torch.int8) + pack_to_int32(weight) + + +class TestAscendW4A16FusedMoEMethod(TestBase): + experts = 8 + input_size = 32 + output_size = 128 + group_size = 32 + + @patch("vllm_ascend.quantization.methods.w4a16.get_ascend_config") + @patch("vllm_ascend.quantization.methods.w4a16.get_current_vllm_config") + def setUp(self, mock_get_current_vllm_config, mock_get_ascend_config): + mock_ascend_config = Mock() + mock_ascend_config.eplb_config.dynamic_eplb = False + mock_ascend_config.eplb_config.expert_map_record_path = None + 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, + }) + mock_get_current_vllm_config.return_value = mock_vllm_config + + self.quant_method = AscendW4A16FusedMoEMethod() + + def test_init(self): + self.assertTrue(self.quant_method.transpose_weight) + self.assertEqual(self.quant_method.num_bits, 4) + self.assertEqual(self.quant_method.pack_factor, 8) + self.assertEqual(self.quant_method.group_size, self.group_size) + self.assertFalse(self.quant_method.dynamic_eplb) + + def test_get_weight(self): + param_dict = self.quant_method.get_weight(self.experts, + self.input_size, + self.output_size, + torch.bfloat16) + + self.assertEqual(param_dict["w13_weight_packed"].dtype, torch.int32) + expected_w13_shape = (self.experts, 2 * self.input_size, + self.output_size // + self.quant_method.pack_factor) + self.assertEqual(param_dict["w13_weight_packed"].shape, + expected_w13_shape) + + self.assertEqual(param_dict["w2_weight_packed"].dtype, torch.int32) + expected_w2_shape = (self.experts, self.output_size, + self.input_size // self.quant_method.pack_factor) + self.assertEqual(param_dict["w2_weight_packed"].shape, + expected_w2_shape) + + def test_get_dynamic_quant_param(self): + 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.bfloat16) + expected_w13_scale_shape = (self.experts, 2 * self.input_size, + self.output_size // self.group_size) + self.assertEqual(param_dict["w13_weight_scale"].shape, + expected_w13_scale_shape) + + self.assertEqual(param_dict["w2_weight_scale"].dtype, torch.bfloat16) + expected_w2_scale_shape = (self.experts, self.output_size, + self.input_size // self.group_size) + self.assertEqual(param_dict["w2_weight_scale"].shape, + expected_w2_scale_shape) + + self.assertEqual(param_dict["w13_weight_shape"].dtype, torch.int32) + self.assertEqual(param_dict["w13_weight_shape"].shape, + (self.experts, 2)) + + self.assertEqual(param_dict["w2_weight_shape"].dtype, torch.int32) + self.assertEqual(param_dict["w2_weight_shape"].shape, + (self.experts, 2)) + + self.assertEqual(param_dict["w13_weight_offset"].dtype, torch.bfloat16) + self.assertEqual(param_dict["w13_weight_offset"].shape, + expected_w13_scale_shape) + + self.assertEqual(param_dict["w2_weight_offset"].dtype, torch.bfloat16) + self.assertEqual(param_dict["w2_weight_offset"].shape, + expected_w2_scale_shape) + + def build_layer(self): + """Build a mock layer for testing""" + layer = torch.nn.Module() + + w13_shape = (self.experts, 2 * self.input_size, + self.output_size // self.quant_method.pack_factor) + w2_shape = (self.experts, self.output_size, + self.input_size // self.quant_method.pack_factor) + + layer.w13_weight_packed = torch.nn.Parameter(torch.randint( + -100, 100, w13_shape, dtype=torch.int32), + requires_grad=False) + layer.w2_weight_packed = torch.nn.Parameter(torch.randint( + -100, 100, w2_shape, dtype=torch.int32), + requires_grad=False) + + w13_scale_shape = (self.experts, 2 * self.input_size, + self.output_size // self.group_size) + w2_scale_shape = (self.experts, self.output_size, + self.input_size // self.group_size) + + layer.w13_weight_scale = torch.nn.Parameter(torch.ones( + w13_scale_shape, dtype=torch.bfloat16), + requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter(torch.ones( + w2_scale_shape, dtype=torch.bfloat16), + requires_grad=False) + + layer.w13_weight_offset = torch.nn.Parameter(torch.zeros( + w13_scale_shape, dtype=torch.bfloat16), + requires_grad=False) + layer.w2_weight_offset = torch.nn.Parameter(torch.zeros( + w2_scale_shape, dtype=torch.bfloat16), + requires_grad=False) + + layer.w13_weight_shape = torch.nn.Parameter(torch.tensor( + [[2 * self.input_size, self.output_size]] * self.experts, + dtype=torch.int32), + requires_grad=False) + layer.w2_weight_shape = torch.nn.Parameter(torch.tensor( + [[self.output_size, self.input_size]] * self.experts, + dtype=torch.int32), + requires_grad=False) + + return layer + + @patch( + "vllm_ascend.quantization.methods.w4a16.torch_npu.npu_convert_weight_to_int4pack" + ) + def test_process_weights_after_loading_with_transpose( + self, mock_npu_convert_weight_to_int4pack): + + def mock_convert_weight(weight): + new_shape = list(weight.shape) + new_shape[-1] = new_shape[-1] // 8 + return torch.zeros(new_shape, dtype=torch.int32) + + mock_npu_convert_weight_to_int4pack.side_effect = mock_convert_weight + + layer = self.build_layer() + self.quant_method.transpose_weight = True + + self.quant_method.process_weights_after_loading(layer) + + self.assertEqual(layer.w13_weight_packed.data.shape, + torch.Size([8, 128, 8])) + self.assertEqual(layer.w2_weight_packed.data.shape, + torch.Size([8, 32, 16])) + + self.assertEqual(layer.w13_weight_scale.data.shape, + torch.Size([8, 4, 64])) + self.assertEqual(layer.w2_weight_scale.data.shape, + torch.Size([8, 1, 128])) + self.assertEqual(layer.w13_weight_offset.data.shape, + torch.Size([8, 4, 64])) + self.assertEqual(layer.w2_weight_offset.data.shape, + torch.Size([8, 1, 128])) + + self.assertTrue(layer.w13_weight_scale.data.is_contiguous()) + self.assertTrue(layer.w2_weight_scale.data.is_contiguous()) + self.assertTrue(layer.w13_weight_offset.data.is_contiguous()) + self.assertTrue(layer.w2_weight_offset.data.is_contiguous()) + + def test_process_weights_after_loading_without_transpose(self): + layer = self.build_layer() + self.quant_method.transpose_weight = False + + original_w13_data = layer.w13_weight_packed.data.clone() + original_w2_data = layer.w2_weight_packed.data.clone() + + self.quant_method.process_weights_after_loading(layer) + + self.assertTrue( + torch.equal(layer.w13_weight_packed.data, original_w13_data)) + self.assertTrue( + torch.equal(layer.w2_weight_packed.data, original_w2_data)) + + @patch("vllm_ascend.quantization.methods.w4a16._EXTRA_CTX") + @patch("vllm_ascend.quantization.methods.w4a16.select_experts") + def test_apply_uses_explicit_dispatch_and_mlp_args(self, mock_select_experts, mock_extra_ctx): + tokens = 3 + hidden_size = self.output_size + layer = self.build_layer() + x = torch.randn(tokens, hidden_size, dtype=torch.float32) + router_logits = torch.randn(tokens, self.experts, dtype=torch.float32) + topk_weights = torch.randn(tokens, 2, dtype=torch.float32) + topk_ids = torch.randint(0, self.experts, (tokens, 2), dtype=torch.int64) + mc2_mask = torch.tensor([1, 0, 1], dtype=torch.bool) + pertoken_scale = torch.randn(tokens, dtype=torch.float32) + + mock_select_experts.return_value = (topk_weights, topk_ids) + mock_comm = Mock() + mock_comm.fused_experts.return_value = torch.randn(tokens, hidden_size, dtype=torch.float32) + mock_extra_ctx.moe_comm_method = mock_comm + mock_extra_ctx.moe_comm_type = MoECommType.ALLGATHER + + self.quant_method.apply( + layer=layer, + x=x, + router_logits=router_logits, + top_k=2, + renormalize=True, + global_num_experts=self.experts, + activation="gelu", + apply_router_weight_on_input=True, + mc2_mask=mc2_mask, + pertoken_scale=pertoken_scale, + ) + + mock_select_experts.assert_called_once() + fused_experts_input = mock_comm.fused_experts.call_args.kwargs["fused_experts_input"] + self.assertEqual(fused_experts_input.activation, "gelu") + self.assertTrue(fused_experts_input.routing.apply_router_weight_on_input) + self.assertIs(fused_experts_input.routing.mc2_mask, mc2_mask) + self.assertIs(fused_experts_input.routing.pertoken_scale, pertoken_scale) + + @patch("vllm_ascend.quantization.methods.w4a16._EXTRA_CTX") + @patch("vllm_ascend.quantization.methods.w4a16.select_experts") + def test_apply_router_logits_mismatch_raises(self, mock_select, mock_ctx): + layer = self.build_layer() + x = torch.randn(4, self.output_size, dtype=torch.float32) + router_logits = torch.randn(4, self.experts + 1, dtype=torch.float32) + with self.assertRaises(AssertionError): + self.quant_method.apply(layer, x, router_logits, top_k=2, renormalize=True, global_num_experts=self.experts) diff --git a/tests/ut/quantization/test_w4a4_flatquant_dynamic.py b/tests/ut/quantization/methods/test_w4a4_flatquant.py similarity index 85% rename from tests/ut/quantization/test_w4a4_flatquant_dynamic.py rename to tests/ut/quantization/methods/test_w4a4_flatquant.py index c3f452c4ea57..b42103fbea49 100644 --- a/tests/ut/quantization/test_w4a4_flatquant_dynamic.py +++ b/tests/ut/quantization/methods/test_w4a4_flatquant.py @@ -6,7 +6,7 @@ from vllm_ascend.quantization.methods.w4a4_flatquant import ( AscendW4A4FlatQuantDynamicLinearMethod, get_decompose_dim, - pack_int4_weights) + pack_int4_weights, KRONECKER_QUANT_MAX_BATCH_SIZE, batched_kronecker_quant) class TestW4A4FlatQuantDynamic(unittest.TestCase): @@ -52,6 +52,39 @@ def test_pack_int4_weights_npu_success(self, mock_torch_npu): mock_torch_npu.npu_convert_weight_to_int4pack.assert_called_once() self.assertTrue(torch.equal(result, mock_packed_tensor)) + @patch('vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu') + def test_large_batch_multiple_calls(self, mock_npu): + batch_size = 50000 + x = torch.randn(batch_size, 24, 32) + left_trans = torch.randn(24, 24) + right_trans = torch.randn(32, 32) + num_chunks = batch_size // KRONECKER_QUANT_MAX_BATCH_SIZE + 1 + mock_returns = [ + (torch.randint(0, 255, (KRONECKER_QUANT_MAX_BATCH_SIZE, 24, 4), dtype=torch.int32), + torch.randn(KRONECKER_QUANT_MAX_BATCH_SIZE)) + for _ in range(num_chunks - 1) + ] + last_chunk_size = batch_size - (num_chunks - 1) * KRONECKER_QUANT_MAX_BATCH_SIZE + mock_returns.append( + (torch.randint(0, 255, (last_chunk_size, 24, 4), dtype=torch.int32), torch.randn(last_chunk_size)) + ) + mock_npu.npu_kronecker_quant.side_effect = mock_returns + result_x, result_scale = batched_kronecker_quant(x, left_trans, right_trans, 0.95) + self.assertEqual(mock_npu.npu_kronecker_quant.call_count, num_chunks) + + @patch('vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu') + def test_exact_max_batch_size(self, mock_npu): + batch_size = KRONECKER_QUANT_MAX_BATCH_SIZE + x = torch.randn(batch_size, 24, 32) + left_trans = torch.randn(24, 24) + right_trans = torch.randn(32, 32) + mock_npu.npu_kronecker_quant.return_value = ( + torch.randint(0, 255, (batch_size, 24, 4), dtype=torch.int32), + torch.randn(batch_size, dtype=torch.float32) + ) + result_x, result_scale = batched_kronecker_quant(x, left_trans, right_trans, 0.95) + mock_npu.npu_kronecker_quant.assert_called_once() + ## Test AscendW4A4FlatQuantDynamicLinearMethod Class ## -------------------------------------------------- diff --git a/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py b/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py new file mode 100644 index 000000000000..0da7b5dbb0fc --- /dev/null +++ b/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py @@ -0,0 +1,94 @@ +from unittest.mock import MagicMock, patch +import torch +import torch.nn as nn +from tests.ut.base import TestBase +from vllm_ascend.quantization.methods.w4a4_laos_dynamic import AscendW4A4LaosDynamicLinearMethod + + +class TestAscendW4A4LaosDynamicLinearMethod(TestBase): + + def setUp(self): + self.method = AscendW4A4LaosDynamicLinearMethod() + + def test_init_transpose_weight_true(self): + self.assertTrue(self.method.transpose_weight) + + def test_get_weight_various_sizes(self): + sizes = [(64, 128), (256, 512), (1024, 2048)] + for input_size, output_size in sizes: + result = self.method.get_weight(input_size, output_size, torch.bfloat16) + self.assertEqual(result["weight"].shape, (output_size, input_size)) + self.assertEqual(result["weight"].dtype, torch.int8) + + def test_get_perchannel_param_various_output_sizes(self): + output_sizes = [1, 64, 128, 512] + for output_size in output_sizes: + result = self.method.get_perchannel_param(output_size, torch.bfloat16) + self.assertEqual(result["weight_scale"].shape, (output_size, 1)) + self.assertEqual(result["weight_offset"].shape, (output_size, 1)) + self.assertEqual(result["weight_scale"].dtype, torch.float32) + self.assertEqual(result["weight_offset"].dtype, torch.float32) + + def test_get_pertensor_param_empty(self): + result = self.method.get_pertensor_param(torch.bfloat16) + self.assertEqual(result, {}) + + def test_get_pergroup_param_various_layer_types(self): + for layer_type in ["column", "row", None]: + result = self.method.get_pergroup_param(256, 128, torch.bfloat16, layer_type) + self.assertEqual(result, {}) + + @patch("torch_npu.npu_quant_matmul") + @patch("torch_npu.npu_dynamic_quant") + def test_apply_2d_input(self, mock_dyn_quant, mock_matmul): + mock_dyn_quant.return_value = ( + torch.randint(0, 15, (32, 128), dtype=torch.int32), + torch.randn(32, dtype=torch.float32) + ) + mock_matmul.return_value = torch.randn(32, 256) + layer = MagicMock() + layer.weight = MagicMock(data=torch.randint(-8, 7, (256, 128), dtype=torch.int8)) + layer.weight_scale = MagicMock(data=torch.randn(256, dtype=torch.float32)) + x = torch.randn(32, 128, dtype=torch.bfloat16) + output = self.method.apply(layer, x) + mock_matmul.assert_called_once() + + @patch("torch_npu.npu_quant_matmul") + @patch("torch_npu.npu_dynamic_quant") + def test_apply_with_bias(self, mock_dyn_quant, mock_matmul): + mock_dyn_quant.return_value = ( + torch.randint(0, 15, (32, 128), dtype=torch.int32), + torch.randn(32, dtype=torch.float32) + ) + expected_output = torch.randn(32, 256, dtype=torch.bfloat16) + mock_matmul.return_value = expected_output + layer = MagicMock() + layer.weight = MagicMock(data=torch.randint(-8, 7, (256, 128), dtype=torch.int8)) + layer.weight_scale = MagicMock(data=torch.randn(256, dtype=torch.float32)) + x = torch.randn(32, 128, dtype=torch.bfloat16) + bias = torch.randn(256, dtype=torch.bfloat16) + output = self.method.apply(layer, x, bias) + expected_output = expected_output + bias + self.assertTrue(torch.equal(output, expected_output)) + + @patch("torch_npu.npu_convert_weight_to_int4pack") + def test_process_weights_no_transpose_when_false(self, mock_convert): + mock_convert.return_value = torch.randint(0, 15, (128, 32), dtype=torch.int32) + self.method.transpose_weight = False + layer = nn.Module() + layer.weight = nn.Parameter(torch.randint(-8, 7, (128, 256), dtype=torch.int8), requires_grad=False) + layer.weight_scale = nn.Parameter(torch.randn(128, 1, dtype=torch.float32), requires_grad=False) + self.method.process_weights_after_loading(layer) + self.assertEqual(layer.weight.shape, (128, 32)) + + @patch("torch_npu.npu_convert_weight_to_int4pack") + def test_process_weights_various_input_sizes(self, mock_convert): + for input_size, output_size in [(64, 128), (256, 512)]: + mock_convert.return_value = torch.randint(0, 15, (output_size, input_size // 8), dtype=torch.int32) + layer = nn.Module() + layer.weight = nn.Parameter(torch.randint(-8, 7, (output_size, input_size), dtype=torch.int8), requires_grad=False) + layer.weight_scale = nn.Parameter(torch.randn(output_size, 1, dtype=torch.float32), requires_grad=False) + self.method.process_weights_after_loading(layer) + mock_convert.assert_called() + self.assertEqual(layer.weight_scale.data.dtype, torch.float32) + self.assertEqual(layer.weight.shape, (input_size // 8, output_size)) diff --git a/tests/ut/quantization/methods/test_w4a4_mxfp4.py b/tests/ut/quantization/methods/test_w4a4_mxfp4.py new file mode 100644 index 000000000000..3f3da32ffe26 --- /dev/null +++ b/tests/ut/quantization/methods/test_w4a4_mxfp4.py @@ -0,0 +1,157 @@ +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_vllm_config, create_mock_ascend_config +from vllm_ascend.device.mxfp_compat import FLOAT8_E8M0FNU_DTYPE +from vllm_ascend.quantization.methods.w4a4_mxfp4 import ( + AscendW4A4MXFP4DynamicLinearMethod, + AscendW4A4MXFP4DynamicFusedMoEMethod, +) +from vllm_ascend.quantization.methods.base import QuantType + + +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_model_dtype_is_none(self): + self.assertIsNone(self.scheme.model_dtype) + + 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_get_pertensor_param_returns_empty(self): + result = self.scheme.get_pertensor_param(torch.bfloat16) + self.assertEqual(result, {}) + + def test_get_perchannel_param_returns_empty(self): + result = self.scheme.get_perchannel_param(128, torch.bfloat16) + self.assertEqual(result, {}) + + 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.ensure_mxfp4_linear_available") + @patch("vllm_ascend.quantization.methods.w4a4_mxfp4.get_current_vllm_config") + def test_group_size_from_config(self, mock_vllm_config, mock_ensure): + mock_vllm_config.return_value = create_mock_vllm_config(quant_description={"group_size": 64}) + scheme = AscendW4A4MXFP4DynamicLinearMethod() + self.assertEqual(scheme.group_size, 64) + + @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_ep_group") + @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_ep, mock_ensure): + mock_vllm.return_value = create_mock_vllm_config() + mock_ascend.return_value = create_mock_ascend_config() + mock_ep.return_value = Mock() + mock_ensure.return_value = None + self.scheme = AscendW4A4MXFP4DynamicFusedMoEMethod() + + def test_init(self): + self.assertEqual(self.scheme.quant_type, QuantType.MXFP4) + self.assertIsNone(self.scheme.model_dtype) + + 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.w2_weight.shape, (8, 128, 128)) + self.assertEqual(layer.w13_weight_scale.shape, (8, 2, 256, 2)) + self.assertEqual(layer.w2_weight_scale.shape, (8, 4, 128, 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() + result = self.scheme.apply( + layer, x, router_logits, top_k=2, renormalize=True, + global_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() diff --git a/tests/ut/quantization/test_w4a8_dynamic.py b/tests/ut/quantization/methods/test_w4a8.py similarity index 73% rename from tests/ut/quantization/test_w4a8_dynamic.py rename to tests/ut/quantization/methods/test_w4a8.py index aedfa61acd06..84549c656735 100644 --- a/tests/ut/quantization/test_w4a8_dynamic.py +++ b/tests/ut/quantization/methods/test_w4a8.py @@ -1,10 +1,12 @@ -from unittest.mock import Mock, patch +from unittest.mock import Mock, patch, MagicMock 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): @@ -120,6 +122,16 @@ def test_process_weights_after_loading(self, mock_format_cast, mock_npu, 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() + class TestAscendW4A8DynamicFusedMoEMethod(TestBase): experts = 8 @@ -272,11 +284,7 @@ def test_process_weights_after_loading(self, mock_npu, mock_npu_quantize, mock_npu_format_cast): mock_npu.return_value = torch.Tensor() mock_npu_quantize.return_value = torch.Tensor() - - def func_by_args(weight, num_format): - return weight - - mock_npu_format_cast.side_effect = func_by_args + mock_npu_format_cast.side_effect = identity # old quant version weight layer = self.build_layer(is_new_quant_version=False) self.quant_method.process_weights_after_loading(layer) @@ -304,3 +312,111 @@ def func_by_args(weight, num_format): 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)) + + 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('torch_npu.npu_quantize') + @patch('torch.Tensor.npu') + def test_process_weights_after_loading_compressed_tensors(self, mock_npu, mock_npu_quantize): + mock_npu.return_value = torch.Tensor() + mock_npu_quantize.return_value = torch.Tensor() + + 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.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) + + @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) + 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, + global_num_experts=num_experts, + expert_map=expert_map, + topk_group=None, + num_expert_group=None, + custom_routing_function=None, + 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["global_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.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)) \ No newline at end of file diff --git a/tests/ut/quantization/test_w8a16.py b/tests/ut/quantization/methods/test_w8a16.py similarity index 58% rename from tests/ut/quantization/test_w8a16.py rename to tests/ut/quantization/methods/test_w8a16.py index 3454cbfdeb89..d2ed247ed9e4 100644 --- a/tests/ut/quantization/test_w8a16.py +++ b/tests/ut/quantization/methods/test_w8a16.py @@ -4,6 +4,7 @@ import torch from tests.ut.base import TestBase +from tests.ut.quantization.conftest_quantization import identity from vllm_ascend.quantization.methods.w8a16 import AscendW8A16LinearMethod @@ -13,9 +14,33 @@ def setUp(self): self.method = AscendW8A16LinearMethod() def test_get_weight(self): - weight = self.method.get_weight(10, 20) + sizes = [(64, 128), (256, 512), (1024, 2048), (1, 1)] + for input_size, output_size in sizes: + weight = self.method.get_weight(input_size, output_size) + self.assertEqual(weight['weight'].dtype, torch.int8) + self.assertEqual(weight['weight'].shape, (output_size, input_size)) + self.assertEqual(len(weight), 1) + + weight = self.method.get_weight(256, 128, torch.float16) self.assertEqual(weight['weight'].dtype, torch.int8) - self.assertEqual(weight['weight'].shape, (20, 10)) + + def test_get_per_channel_param(self): + for output_size, dtype in [(128, torch.bfloat16), (256, torch.float16)]: + per_channel_params = self.method.get_perchannel_param(output_size, dtype) + self.assertEqual(per_channel_params['weight_scale'].dtype, dtype) + self.assertEqual(per_channel_params['weight_scale'].shape, (output_size, 1)) + self.assertEqual(per_channel_params['weight_offset'].dtype, dtype) + self.assertEqual(per_channel_params['weight_offset'].shape, (output_size, 1)) + self.assertEqual(len(per_channel_params), 2) + + def test_get_pertensor_param_returns_empty(self): + result = self.method.get_pertensor_param(torch.bfloat16) + self.assertEqual(result, {}) + + def test_get_pergroup_param_returns_empty(self): + result = self.method.get_pergroup_param(256, 128, torch.bfloat16) + self.assertEqual(result, {}) + @patch("torch_npu.npu_weight_quant_batchmatmul") def test_apply_with_x_is_int8(self, mock_npu_weight_quant_batchmatmul): @@ -31,23 +56,24 @@ def test_apply_with_x_is_int8(self, mock_npu_weight_quant_batchmatmul): mock_npu_weight_quant_batchmatmul.return_value = expected_y_output output = self.method.apply(layer, x, bias) - expected_y_output += bias self.assertTrue(torch.equal(output, expected_y_output)) + mock_npu_weight_quant_batchmatmul.assert_called_once() @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "0"}) @patch('torch_npu.npu_format_cast') def test_process_weights_after_loading_with_nz0(self, mock_npu_format_cast): layer = MagicMock() - layer.weight.data = torch.randint(-127, - 128, (128, 256), + layer.weight.data = torch.randint(-128, + 127, (128, 256), dtype=torch.int8) layer.weight_scale.data = torch.randn(128, 1) layer.weight_offset.data = torch.randn(128, 1) - mock_npu_format_cast.return_value = MagicMock + mock_npu_format_cast.side_effect = identity self.method.process_weights_after_loading(layer) + self.assertEqual(layer.weight.data.shape, (256, 128)) self.assertEqual(layer.weight_scale.data.shape, (128, )) self.assertEqual(layer.weight_offset.data.shape, (128, )) mock_npu_format_cast.assert_not_called() @@ -58,15 +84,16 @@ def test_process_weights_after_loading_with_nz1(self, mock_npu_format_cast): layer = MagicMock() - layer.weight.data = torch.randint(-127, - 128, (128, 256), + layer.weight.data = torch.randint(-128, + 127, (128, 256), dtype=torch.int8) layer.weight_scale.data = torch.randn(128, 1) layer.weight_offset.data = torch.randn(128, 1) - mock_npu_format_cast.return_value = MagicMock + mock_npu_format_cast.side_effect = identity self.method.process_weights_after_loading(layer) + self.assertEqual(layer.weight.data.shape, (256, 128)) self.assertEqual(layer.weight_scale.data.shape, (128, )) self.assertEqual(layer.weight_offset.data.shape, (128, )) mock_npu_format_cast.assert_called_once() @@ -77,15 +104,16 @@ def test_process_weights_after_loading_with_nz2(self, mock_npu_format_cast): layer = MagicMock() - layer.weight.data = torch.randint(-127, - 128, (128, 256), + layer.weight.data = torch.randint(-128, + 127, (128, 256), dtype=torch.int8) layer.weight_scale.data = torch.randn(128, 1) layer.weight_offset.data = torch.randn(128, 1) - mock_npu_format_cast.return_value = MagicMock + mock_npu_format_cast.side_effect = identity self.method.process_weights_after_loading(layer) + self.assertEqual(layer.weight.data.shape, (256, 128)) self.assertEqual(layer.weight_scale.data.shape, (128, )) self.assertEqual(layer.weight_offset.data.shape, (128, )) mock_npu_format_cast.assert_called_once() diff --git a/tests/ut/quantization/methods/test_w8a8_dynamic.py b/tests/ut/quantization/methods/test_w8a8_dynamic.py new file mode 100644 index 000000000000..0fa6ed9d08e4 --- /dev/null +++ b/tests/ut/quantization/methods/test_w8a8_dynamic.py @@ -0,0 +1,339 @@ +from unittest.mock import Mock, patch, MagicMock + +import torch + +from tests.ut.base import TestBase +from tests.ut.quantization.conftest_quantization import create_mock_vllm_config, create_mock_ascend_config, \ + create_moe_layer +from vllm_ascend.ascend_forward_context import MoECommType +from vllm_ascend.quantization.methods.w8a8_dynamic import ( + AscendW8A8DynamicLinearMethod, + AscendW8A8DynamicFusedMoEMethod, + scale_from_float_to_int64, +) +from vllm_ascend.quantization.quant_type import QuantType + + +class TestScaleFromFloatToInt64(TestBase): + + def test_scale_conversion_basic(self): + scale = torch.tensor([0.5], dtype=torch.float32) + with patch.object(scale, 'cpu', return_value=scale): + with patch.object(scale, 'to', return_value=scale): + result = scale_from_float_to_int64(scale) + self.assertEqual(result.dtype, torch.int64) + + def test_scale_conversion_preserves_device(self): + scale = torch.tensor([0.5, 1.0], dtype=torch.float32) + with patch.object(scale, 'cpu', return_value=scale): + with patch.object(scale, 'to', return_value=scale): + result = scale_from_float_to_int64(scale) + self.assertEqual(result.shape, (2,)) + + +class TestAscendW8A8DynamicLinearMethod(TestBase): + + def setUp(self): + self.method = AscendW8A8DynamicLinearMethod() + + def test_get_weight_various_sizes(self): + sizes = [(64, 128), (256, 512), (1024, 2048)] + for input_size, output_size in sizes: + weight = self.method.get_weight(input_size, output_size, torch.bfloat16) + self.assertEqual(weight['weight'].dtype, torch.int8) + self.assertEqual(weight['weight'].shape, (output_size, input_size)) + + def test_get_perchannel_param_dtype_variations(self): + dtypes = [torch.bfloat16, torch.float16] + for dtype in dtypes: + params = self.method.get_perchannel_param(128, dtype) + self.assertEqual(params['weight_scale'].dtype, dtype) + self.assertEqual(params['weight_offset'].dtype, dtype) + self.assertEqual(params['weight_scale'].shape, (128, 1)) + self.assertEqual(params['weight_offset'].shape, (128, 1)) + + @patch("torch_npu.npu_quant_matmul") + @patch("torch_npu.npu_dynamic_quant") + def test_apply_2d_input(self, mock_dyn_quant, mock_matmul): + mock_dyn_quant.return_value = ( + torch.randint(-128, 127, (32, 128), dtype=torch.int8), + torch.randn(32, dtype=torch.float32) + ) + mock_matmul.return_value = torch.randn(32, 256) + layer = MagicMock() + layer.weight = torch.randint(-128, 127, (128, 256), dtype=torch.int8) + layer.weight_scale = torch.randn(256, dtype=torch.float32) + x = torch.randn(32, 128, dtype=torch.bfloat16) + output = self.method.apply(layer, x) + mock_dyn_quant.assert_called_once() + mock_matmul.assert_called_once() + + @patch("torch_npu.npu_quant_matmul") + @patch("torch_npu.npu_dynamic_quant") + def test_apply_3d_input_with_squeeze(self, mock_dyn_quant, mock_matmul): + mock_dyn_quant.return_value = ( + torch.randint(-128, 127, (32, 1, 128), dtype=torch.int8), + torch.randn(32, 1, dtype=torch.float32) + ) + mock_matmul.return_value = torch.randn(32, 1, 256) + layer = MagicMock() + layer.weight = torch.randint(-128, 127, (128, 256), dtype=torch.int8) + layer.weight_scale = torch.randn(256, dtype=torch.float32) + x = torch.randn(32, 1, 128, dtype=torch.bfloat16) + output = self.method.apply(layer, x) + self.assertEqual(output.shape, (32, 1, 1, 256)) + + def test_process_weights_after_loading(self): + layer = MagicMock() + layer.weight.data = torch.randint(-128, 127, (128, 256), dtype=torch.int8) + layer.weight_scale.data=torch.randn(256, 1, dtype=torch.bfloat16) + layer.weight_offset.data=torch.randn(256, 1, dtype=torch.bfloat16) + with patch('vllm_ascend.quantization.methods.w8a8_dynamic.maybe_trans_nz', side_effect=lambda x: x): + self.method.process_weights_after_loading(layer) + self.assertEqual(layer.weight_scale_fp32.dtype, torch.float32) + self.assertEqual(layer.weight_scale.data.shape, (256,)) + self.assertEqual(layer.weight_offset.data.shape, (256,)) + self.assertEqual(layer.weight.data.shape, (256, 128)) + + +class TestAscendW8A8FusedMoEMethod(TestBase): + num_experts = 8 + hidden_size = 128 + intermediate_size = 128 + + @patch("torch.distributed.get_rank") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_mc2_group") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ascend_config") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ep_group") + def setUp(self, mock_ep, mock_ascend, mock_mc2, mock_rank): + with patch('vllm_ascend.quantization.methods.w8a8_dynamic.get_current_vllm_config') as mock_vllm: + mock_vllm.return_value = create_mock_vllm_config() + mock_ep.return_value = Mock() + mock_ascend.return_value = create_mock_ascend_config() + mock_mc2.return_value = MagicMock(device_group=Mock( + _get_backend=Mock(return_value=Mock(get_hccl_comm_name=Mock(return_value="test_comm"))))) + mock_rank.return_value = 0 + self.quant_method = AscendW8A8DynamicFusedMoEMethod() + + def test_quant_type_is_w8a8(self): + self.assertEqual(self.quant_method.quant_type, QuantType.W8A8) + + def test_get_weight_various_expert_counts(self): + expert_counts = [4, 8, 16, 32] + for num_experts in expert_counts: + param_dict = self.quant_method.get_weight(num_experts, self.intermediate_size, self.hidden_size, + torch.bfloat16) + self.assertEqual(param_dict["w13_weight"].shape[0], num_experts) + self.assertEqual(param_dict["w2_weight"].shape[0], num_experts) + + def test_get_dynamic_quant_param_various_sizes(self): + param_dict = self.quant_method.get_dynamic_quant_param(self.num_experts, self.intermediate_size, + self.hidden_size, torch.bfloat16) + self.assertEqual(param_dict["w13_weight_scale"].dtype, torch.bfloat16) + self.assertEqual(param_dict["w13_weight_offset"].shape, (self.num_experts, 2 * self.intermediate_size, 1)) + self.assertEqual(param_dict["w2_weight_scale"].dtype, torch.bfloat16) + self.assertEqual(param_dict["w2_weight_offset"].shape, (self.num_experts, self.hidden_size, 1)) + + @patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts") + def test_apply_uses_explicit_dispatch_and_mlp_args(self, mock_select_experts, mock_extra_ctx): + tokens = 4 + hidden_size = self.hidden_size + layer = torch.nn.Module() + layer.w13_weight = torch.randint( + -8, + 8, + (self.num_experts, 2 * self.intermediate_size, hidden_size), + dtype=torch.int8, + ) + layer.w2_weight = torch.randint( + -8, + 8, + (self.num_experts, hidden_size, self.intermediate_size), + dtype=torch.int8, + ) + layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size, dtype=torch.float32) + layer.w2_weight_scale = torch.ones(self.num_experts, hidden_size, dtype=torch.float32) + + x = torch.randn(tokens, hidden_size, dtype=torch.float32) + router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32) + topk_weights = torch.randn(tokens, 2, dtype=torch.float32) + topk_ids = torch.randint(0, self.num_experts, (tokens, 2), dtype=torch.int64) + mc2_mask = torch.tensor([1, 0, 1, 0], dtype=torch.bool) + pertoken_scale = torch.randn(tokens, dtype=torch.float32) + + mock_select_experts.return_value = (topk_weights, topk_ids) + mock_comm = Mock() + mock_comm.fused_experts.return_value = torch.randn(tokens, hidden_size, dtype=torch.float32) + mock_extra_ctx.moe_comm_method = mock_comm + mock_extra_ctx.moe_comm_type = MoECommType.ALLGATHER + self.quant_method.multistream_overlap_gate = False + self.quant_method.in_dtype = torch.float32 + + self.quant_method.apply( + layer=layer, + x=x, + router_logits=router_logits, + top_k=2, + renormalize=True, + global_num_experts=self.num_experts, + activation="gelu", + apply_router_weight_on_input=True, + mc2_mask=mc2_mask, + pertoken_scale=pertoken_scale, + ) + + fused_experts_input = mock_comm.fused_experts.call_args.kwargs["fused_experts_input"] + self.assertEqual(fused_experts_input.activation, "gelu") + self.assertTrue(fused_experts_input.routing.apply_router_weight_on_input) + self.assertIs(fused_experts_input.routing.mc2_mask, mc2_mask) + self.assertIs(fused_experts_input.routing.pertoken_scale, pertoken_scale) + self.assertIs(fused_experts_input.topk_weights, topk_weights) + self.assertIs(fused_experts_input.topk_ids, topk_ids) + + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_flash_common3_context") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts") + def test_apply_overlap_gate_uses_fc3_context( + self, + mock_select_experts, + mock_extra_ctx, + mock_get_flash_common3_context, + ): + tokens = 4 + hidden_size = self.hidden_size + layer = torch.nn.Module() + layer.w13_weight = torch.randint( + -8, + 8, + (self.num_experts, 2 * self.intermediate_size, hidden_size), + dtype=torch.int8, + ) + layer.w2_weight = torch.randint( + -8, + 8, + (self.num_experts, hidden_size, self.intermediate_size), + dtype=torch.int8, + ) + layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size, dtype=torch.float32) + layer.w2_weight_scale = torch.ones(self.num_experts, hidden_size, dtype=torch.float32) + + x = torch.randn(tokens, hidden_size, dtype=torch.float32) + router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32) + topk_weights = torch.randn(tokens, 2, dtype=torch.float32) + topk_ids = torch.randint(0, self.num_experts, (tokens, 2), dtype=torch.int64) + mc2_mask = torch.tensor([1, 0, 1, 0], dtype=torch.bool) + pertoken_scale = torch.randn(tokens, dtype=torch.float32) + + self.quant_method.multistream_overlap_gate = True + self.quant_method.in_dtype = torch.float32 + mock_get_flash_common3_context.return_value = Mock(topk_weights=topk_weights, topk_ids=topk_ids) + + mock_comm = Mock() + mock_comm.fused_experts.return_value = torch.randn(tokens, hidden_size, dtype=torch.float32) + mock_extra_ctx.moe_comm_method = mock_comm + mock_extra_ctx.moe_comm_type = MoECommType.ALLGATHER + + self.quant_method.apply( + layer=layer, + x=x, + router_logits=router_logits, + top_k=2, + renormalize=True, + global_num_experts=self.num_experts, + activation="gelu", + apply_router_weight_on_input=True, + mc2_mask=mc2_mask, + pertoken_scale=pertoken_scale, + ) + + mock_select_experts.assert_not_called() + fused_experts_input = mock_comm.fused_experts.call_args.kwargs["fused_experts_input"] + self.assertEqual(fused_experts_input.activation, "gelu") + self.assertTrue(fused_experts_input.routing.apply_router_weight_on_input) + self.assertIs(fused_experts_input.routing.mc2_mask, mc2_mask) + self.assertIs(fused_experts_input.routing.pertoken_scale, pertoken_scale) + self.assertIs(fused_experts_input.topk_weights, topk_weights) + self.assertIs(fused_experts_input.topk_ids, topk_ids) + + @patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.zero_experts_compute") + def test_apply_with_zero_experts(self, mock_zero, mock_select, mock_ctx): + tokens = 4 + layer = MagicMock() + layer.w13_weight = torch.randint(-8, 8, (self.num_experts, 2 * self.intermediate_size, self.hidden_size), + dtype=torch.int8) + layer.w2_weight = torch.randint(-8, 8, (self.num_experts, self.hidden_size, self.intermediate_size), + dtype=torch.int8) + layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size) + layer.w2_weight_scale = torch.ones(self.num_experts, self.hidden_size) + layer.zero_expert_num = 2 + layer.zero_expert_type = "shared" + layer.n_shared_experts = 0 + layer.mix_placement = False + x = torch.randn(tokens, self.hidden_size, dtype=torch.float32) + 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_zero.return_value = (topk_ids, topk_weights, torch.randn(tokens, self.hidden_size)) + 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.quant_method.in_dtype = torch.float32 + self.quant_method.apply(layer, x, router_logits, top_k=2, renormalize=True, global_num_experts=self.num_experts) + + @patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts") + def test_apply_with_enable_force_load_balance(self, mock_select, mock_ctx): + tokens = 4 + layer = MagicMock() + layer.w13_weight = torch.randint(-8, 8, (self.num_experts, 2 * self.intermediate_size, self.hidden_size), + dtype=torch.int8) + layer.w2_weight = torch.randint(-8, 8, (self.num_experts, self.hidden_size, self.intermediate_size), + dtype=torch.int8) + layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size) + layer.w2_weight_scale = torch.ones(self.num_experts, self.hidden_size) + layer.zero_expert_num = 0 + layer.zero_expert_type = None + layer.n_shared_experts = 0 + layer.mix_placement = False + x = torch.randn(tokens, self.hidden_size, dtype=torch.float32) + 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.quant_method.in_dtype = torch.float32 + self.quant_method.apply(layer, x, router_logits, top_k=2, renormalize=True, global_num_experts=self.num_experts, + enable_force_load_balance=True) + + @patch('torch_npu.npu_format_cast') + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.envs_ascend") + def test_process_weights_with_fused_mc2(self, mock_envs, mock_format_cast): + mock_envs.VLLM_ASCEND_ENABLE_FUSED_MC2 = 1 + mock_format_cast.return_value = torch.randint(-8, 8, + (self.num_experts, self.hidden_size, 2 * self.intermediate_size), + dtype=torch.int8) + layer = create_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size) + self.quant_method.process_weights_after_loading(layer) + self.assertTrue(hasattr(layer, 'w13_weight_scale_fp32')) + + @patch('torch_npu.npu_format_cast') + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.envs_ascend") + def test_process_weights_with_dynamic_eplb(self, mock_envs, mock_format_cast): + mock_envs.VLLM_ASCEND_ENABLE_FUSED_MC2 = 0 + self.quant_method.dynamic_eplb = True + mock_format_cast.return_value = torch.randint(-8, 8, + (self.num_experts, self.hidden_size, 2 * self.intermediate_size), + dtype=torch.int8) + layer = create_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size) + self.quant_method.process_weights_after_loading(layer) + self.assertTrue(hasattr(layer, 'w13_weight_list')) diff --git a/tests/ut/quantization/methods/test_w8a8_mxfp8.py b/tests/ut/quantization/methods/test_w8a8_mxfp8.py new file mode 100644 index 000000000000..f3a740bcfd25 --- /dev/null +++ b/tests/ut/quantization/methods/test_w8a8_mxfp8.py @@ -0,0 +1,220 @@ +from unittest.mock import Mock, patch + +import torch +import torch.nn as nn + +from tests.ut.base import TestBase +from tests.ut.quantization.conftest_quantization import ( + create_mock_vllm_config, + create_mock_ascend_config, + create_mxfp_moe_layer, +) +from vllm_ascend.quantization.methods.base import QuantType +from vllm_ascend.quantization.methods.w8a8_mxfp8 import ( + AscendW8A8MXFP8DynamicLinearMethod, + AscendW8A8MXFP8DynamicFusedMoEMethod, +) + + +class TestAscendW8A8MXFP8LinearMethod(TestBase): + + @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.ensure_mxfp8_linear_available") + @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.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 = AscendW8A8MXFP8DynamicLinearMethod() + + def test_init(self): + self.assertIsNone(self.scheme.model_dtype) + self.assertEqual(self.scheme.group_size, 32) + + def test_get_weight_float8_dtype(self): + result = self.scheme.get_weight(256, 128, torch.bfloat16) + self.assertEqual(result["weight"].dtype, torch.float8_e4m3fn) + self.assertEqual(len(result), 1) + + def test_get_weight_various_input_sizes(self): + sizes = [(128, 64), (512, 256), (1024, 512)] + for input_size, output_size in sizes: + result = self.scheme.get_weight(input_size, output_size, torch.bfloat16) + self.assertEqual(result["weight"].shape, (output_size, input_size)) + + def test_get_pergroup_param_group_size_variations(self): + group_sizes = [16, 32, 64, 128] + 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_stores_original_shapes(self): + layer = nn.Module() + layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), 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.assertTrue(hasattr(layer, "_mxfp8_original_shapes")) + self.assertEqual(layer._mxfp8_original_shapes["weight"], (128, 256)) + self.assertTrue(layer._mxfp8_transformed) + + def test_process_weights_double_call_no_change(self): + layer = nn.Module() + layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), 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) + weight_after_first = layer.weight.data.clone() + self.scheme.process_weights_after_loading(layer) + self.assertTrue(torch.equal(layer.weight.data, weight_after_first)) + + def test_restore_without_mxfp8_transformed(self): + layer = nn.Module() + layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), requires_grad=False) + original_weight_shape = layer.weight.shape + layer._mxfp8_transformed = False + self.scheme.restore_weights_for_rl_loading(layer) + self.assertEqual(layer.weight.shape, original_weight_shape) + + def test_restore_without_original_shapes_raises(self): + layer = nn.Module() + layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), requires_grad=False) + layer.weight_scale = nn.Parameter(torch.randint(0, 255, (128, 8), dtype=torch.uint8), requires_grad=False) + layer._mxfp8_transformed = True + with self.assertRaises(RuntimeError): + self.scheme.restore_weights_for_rl_loading(layer) + + def test_restore_after_process_returns_original_shape(self): + layer = nn.Module() + layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), requires_grad=False) + layer.weight_scale = nn.Parameter(torch.randint(0, 255, (128, 8), dtype=torch.uint8), requires_grad=False) + original_weight_shape = layer.weight.shape + original_scale_shape = layer.weight_scale.shape + self.scheme.process_weights_after_loading(layer) + self.scheme.restore_weights_for_rl_loading(layer) + self.assertEqual(layer.weight.shape, original_weight_shape) + self.assertEqual(layer.weight_scale.shape, original_scale_shape) + self.assertFalse(layer._mxfp8_transformed) + + def test_scale_shape_transformation(self): + layer = nn.Module() + layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), 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_scale.shape, (4, 128, 2)) + + @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.torch_npu") + def test_apply(self, mock_torch_npu): + from vllm_ascend.device.mxfp_compat import FLOAT8_E8M0FNU_DTYPE + dynamic_scale = torch.randint(0, 255, (32, 8), dtype=torch.uint8) + mock_torch_npu.npu_dynamic_mx_quant.return_value = ( + torch.randint(0, 255, (32, 256), dtype=torch.uint8), + dynamic_scale + ) + mock_torch_npu.npu_quant_matmul.return_value = torch.randn(32, 128, dtype=torch.float16) + layer = nn.Module() + layer.weight = nn.Parameter(torch.randn(256, 128).to(torch.float8_e4m3fn), requires_grad=False) + layer.weight_scale = nn.Parameter(torch.randint(0, 255, (4, 128, 2), dtype=torch.uint8), requires_grad=False) + x = torch.randn(32, 1, 256, dtype=torch.float16) + bias = torch.randn(128, dtype=torch.float16) + output = self.scheme.apply(layer, x, bias) + self.assertEqual(output.shape, (32, 1, 128)) + call_kwargs = mock_torch_npu.npu_quant_matmul.call_args.kwargs + self.assertEqual(call_kwargs['bias'].dtype, torch.float32) + self.assertEqual(call_kwargs['group_sizes'], [1, 1, self.scheme.group_size]) + self.assertEqual(call_kwargs['scale_dtype'], FLOAT8_E8M0FNU_DTYPE) + self.assertEqual(call_kwargs['pertoken_scale_dtype'], FLOAT8_E8M0FNU_DTYPE) + self.assertEqual(call_kwargs['output_dtype'], torch.float16) + + +class TestAscendW8A8MXFP8MoEMethod(TestBase): + + num_experts = 8 + hidden_size = 128 + intermediate_size = 256 + + @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.ensure_mxfp8_moe_available") + @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.get_ep_group") + @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.get_current_vllm_config") + @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.get_ascend_config") + def setUp(self, mock_ascend, mock_vllm, mock_ep, mock_ensure): + mock_vllm.return_value = create_mock_vllm_config() + mock_ascend.return_value = create_mock_ascend_config() + mock_ep.return_value = Mock() + mock_ensure.return_value = None + self.scheme = AscendW8A8MXFP8DynamicFusedMoEMethod() + + def test_quant_type_is_mxfp8(self): + self.assertEqual(self.scheme.quant_type, QuantType.MXFP8) + + def test_model_dtype_is_none(self): + self.assertIsNone(self.scheme.model_dtype) + + 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.float8_e4m3fn) + self.assertEqual(result["w2_weight"].dtype, torch.float8_e4m3fn) + + def test_get_weight_various_expert_counts(self): + for num_experts in [4, 8, 16]: + result = self.scheme.get_weight(num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16) + self.assertEqual(result["w13_weight"].shape[0], num_experts) + + def test_get_dynamic_quant_param_dtype_uint8(self): + result = self.scheme.get_dynamic_quant_param(self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16) + self.assertEqual(result["w13_weight_scale"].dtype, torch.uint8) + self.assertEqual(result["w2_weight_scale"].dtype, torch.uint8) + + def test_process_weights_stores_original_shapes(self): + layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + self.scheme.process_weights_after_loading(layer) + self.assertTrue(hasattr(layer, "_mxfp8_original_shapes")) + self.assertIn("w13_weight", layer._mxfp8_original_shapes) + + def test_process_weights_double_call_idempotent(self): + layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + self.scheme.process_weights_after_loading(layer) + weight_after = layer.w13_weight.data.clone() + self.scheme.process_weights_after_loading(layer) + self.assertTrue(torch.equal(layer.w13_weight.data, weight_after)) + + def test_restore_weights_for_rl_loading(self): + layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + original_w13_shape = layer.w13_weight.shape + self.scheme.process_weights_after_loading(layer) + self.assertNotEqual(layer.w13_weight.shape, original_w13_shape) + self.scheme.restore_weights_for_rl_loading(layer) + self.assertEqual(layer.w13_weight.shape, original_w13_shape) + + def test_restore_without_transform_is_noop(self): + layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + original_weight = layer.w13_weight.data.clone() + self.scheme.restore_weights_for_rl_loading(layer) + self.assertTrue(torch.equal(layer.w13_weight.data, original_weight)) + + def test_process_weights_transposes_weights(self): + layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + original_shape = layer.w13_weight.shape + self.scheme.process_weights_after_loading(layer) + self.assertEqual(layer.w13_weight.shape, (original_shape[0], original_shape[2], original_shape[1])) + + @patch("vllm_ascend.quantization.methods.w8a8_mxfp8._EXTRA_CTX") + @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.select_experts") + def test_apply_full_params(self, mock_select, mock_ctx): + tokens = 4 + layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + self.scheme.process_weights_after_loading(layer) + 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() + result = self.scheme.apply( + layer, x, router_logits, top_k=2, renormalize=True, + global_num_experts=self.num_experts, + activation="silu", pertoken_scale=torch.randn(tokens) + ) + mock_select.assert_called_once() + mock_comm.fused_experts.assert_called_once() \ No newline at end of file diff --git a/tests/ut/quantization/methods/test_w8a8_pdmix.py b/tests/ut/quantization/methods/test_w8a8_pdmix.py new file mode 100644 index 000000000000..723b8ec39b3a --- /dev/null +++ b/tests/ut/quantization/methods/test_w8a8_pdmix.py @@ -0,0 +1,164 @@ +from unittest.mock import MagicMock, patch +import torch +from tests.ut.base import TestBase +from tests.ut.quantization.conftest_quantization import create_mock_vllm_config +from vllm_ascend.quantization.methods import AscendW8A8LinearMethod, AscendW8A8DynamicLinearMethod, \ + AscendW8A8PDMixLinearMethod, AscendW8A8PDMixFusedMoeMethod + + +class TestAscendW8A8PDMixLinearScheme(TestBase): + def setUp(self): + self.method = AscendW8A8LinearMethod() + + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") + def test_init_creates_static_and_dynamic_methods(self, mock_vllm_config): + mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None) + scheme = AscendW8A8PDMixLinearMethod() + self.assertTrue(isinstance(scheme._static_method, AscendW8A8LinearMethod)) + self.assertTrue(isinstance(scheme._dynamic_method, AscendW8A8DynamicLinearMethod)) + self.assertFalse(scheme._is_kv_consumer) + + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") + def test_get_weight_delegates_to_static(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls): + mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None) + mock_dynamic_instance = MagicMock() + mock_dynamic_cls.return_value = mock_dynamic_instance + mock_static_instance = MagicMock() + mock_static_instance.get_weight.return_value = {"weight": torch.empty(128, 256, dtype=torch.int8)} + mock_static_cls.return_value = mock_static_instance + scheme = AscendW8A8PDMixLinearMethod() + for input_size, output_size in [(64, 128), (256, 512), (1024, 2048)]: + result = scheme.get_weight(input_size, output_size, torch.bfloat16) + mock_static_instance.get_weight.assert_called_with(input_size, output_size, torch.bfloat16) + mock_dynamic_instance.get_weight.assert_not_called() + + + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") + def test_get_pertensor_param_delegates_to_static(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls): + mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None) + mock_dynamic_instance = MagicMock() + mock_dynamic_cls.return_value = mock_dynamic_instance + mock_static_instance = MagicMock() + mock_static_instance.get_pertensor_param.return_value = {} + mock_static_cls.return_value = mock_static_instance + scheme = AscendW8A8PDMixLinearMethod() + result = scheme.get_pertensor_param(torch.bfloat16) + mock_static_instance.get_pertensor_param.assert_called_once_with(torch.bfloat16) + mock_dynamic_instance.get_pertensor_param.assert_not_called() + + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") + def test_get_perchannel_param_delegates_to_static(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls): + mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None) + mock_dynamic_instance = MagicMock() + mock_dynamic_cls.return_value = mock_dynamic_instance + mock_static_instance = MagicMock() + mock_static_instance.get_perchannel_param.return_value = {} + mock_static_cls.return_value = mock_static_instance + scheme = AscendW8A8PDMixLinearMethod() + result = scheme.get_perchannel_param(128, torch.bfloat16) + mock_static_instance.get_perchannel_param.assert_called_once_with(128, torch.bfloat16) + mock_dynamic_instance.get_perchannel_param.assert_not_called() + + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") + def test_apply_uses_static_for_kv_consumer(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls): + mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None) + mock_static_instance = MagicMock() + mock_static_instance.apply.return_value = torch.randn(4, 128) + mock_static_cls.return_value = mock_static_instance + mock_dynamic_instance = MagicMock() + mock_dynamic_cls.return_value = mock_dynamic_instance + scheme = AscendW8A8PDMixLinearMethod() + layer = MagicMock() + layer.is_kv_consumer = True + x = torch.randn(4, 256) + scheme.apply(layer, x) + mock_static_instance.apply.assert_called_once() + mock_dynamic_instance.apply.assert_not_called() + + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") + def test_apply_uses_dynamic_for_non_kv_consumer(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls): + mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None) + mock_dynamic_instance = MagicMock() + mock_dynamic_instance.apply.return_value = torch.randn(4, 128) + mock_dynamic_cls.return_value = mock_dynamic_instance + mock_static_instance = MagicMock() + mock_static_cls.return_value = mock_static_instance + scheme = AscendW8A8PDMixLinearMethod() + layer = MagicMock() + layer.is_kv_consumer = False + x = torch.randn(4, 256) + scheme.apply(layer, x) + mock_dynamic_instance.apply.assert_called_once() + mock_static_instance.apply.assert_not_called() + + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") + def test_process_weights_after_loading_sets_is_kv_consumer(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls): + mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None) + mock_static_instance = MagicMock() + mock_static_cls.return_value = mock_static_instance + mock_dynamic_instance = MagicMock() + mock_dynamic_cls.return_value = mock_dynamic_instance + scheme = AscendW8A8PDMixLinearMethod() + layer = MagicMock() + layer.weight_scale = MagicMock(data=torch.randn(128, 1, dtype=torch.bfloat16)) + scheme.process_weights_after_loading(layer) + mock_static_instance.process_weights_after_loading.assert_called_once_with(layer) + mock_dynamic_instance.process_weights_after_loading.assert_not_called() + self.assertFalse(layer.is_kv_consumer) + + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") + @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") + def test_is_kv_consumer_when_kv_transfer_config_set(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls): + mock_kv_config = MagicMock() + mock_kv_config.is_kv_consumer = True + mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=mock_kv_config) + mock_static_instance = MagicMock() + mock_static_cls.return_value = mock_static_instance + scheme = AscendW8A8PDMixLinearMethod() + layer = MagicMock() + layer.weight_scale = MagicMock(data=torch.randn(128, 1, dtype=torch.bfloat16)) + scheme.process_weights_after_loading(layer) + self.assertTrue(layer.is_kv_consumer) + + +class TestAscendW8A8PDMixMoEScheme(TestBase): + + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_mc2_group") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ep_group") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_current_vllm_config") + @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ascend_config") + def test_get_dynamic_quant_param(self, mock_ascend, mock_vllm, mock_ep, mock_mc2): + mock_mc2.side_effect = AttributeError() + mock_vllm.return_value = create_mock_vllm_config() + mock_ascend.return_value = MagicMock(eplb_config=MagicMock(dynamic_eplb=False)) + scheme = AscendW8A8PDMixFusedMoeMethod() + num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype = 8, 256, 128, torch.bfloat16 + result = scheme.get_dynamic_quant_param(num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype) + # test adds extra params + self.assertEqual(result["w2_deq_scale"].shape, (num_experts, hidden_sizes)) + self.assertEqual(result["w2_deq_scale"].dtype, torch.float32) + self.assertEqual(result["w13_deq_scale"].shape, (num_experts, 2 * intermediate_size_per_partition)) + self.assertEqual(result["w13_deq_scale"].dtype, torch.float32) + self.assertEqual(result["w2_input_offset"].shape, (num_experts, 1)) + self.assertEqual(result["w2_input_offset"].dtype, torch.int8) + self.assertEqual(result["w13_input_offset"].shape, (num_experts, 1)) + self.assertEqual(result["w13_input_offset"].dtype, torch.int8) + + # test parent params + self.assertIn("w13_weight_scale", result) + self.assertIn("w13_weight_offset", result) + self.assertIn("w2_weight_scale", result) + self.assertIn("w2_weight_offset", result) \ No newline at end of file diff --git a/tests/ut/quantization/test_w8a8.py b/tests/ut/quantization/methods/test_w8a8_static.py similarity index 57% rename from tests/ut/quantization/test_w8a8.py rename to tests/ut/quantization/methods/test_w8a8_static.py index bbed09e6c781..fd4fc273ae5a 100644 --- a/tests/ut/quantization/test_w8a8.py +++ b/tests/ut/quantization/methods/test_w8a8_static.py @@ -4,8 +4,9 @@ import torch from tests.ut.base import TestBase +from tests.ut.quantization.conftest_quantization import identity from vllm_ascend.quantization.methods.w8a8_static import AscendW8A8LinearMethod -from vllm_ascend.utils import AscendDeviceType +from vllm_ascend.utils import AscendDeviceType, COMPRESSED_TENSORS_METHOD class TestAscendW8A8LinearMethod(TestBase): @@ -14,28 +15,39 @@ def setUp(self): self.method = AscendW8A8LinearMethod() def test_get_weight(self): - weight = self.method.get_weight(10, 20) + sizes = [(64, 128), (256, 512), (1024, 2048), (1, 1)] + for input_size, output_size in sizes: + weight = self.method.get_weight(input_size, output_size) + self.assertEqual(weight['weight'].dtype, torch.int8) + self.assertEqual(weight['weight'].shape, (output_size, input_size)) + self.assertEqual(len(weight), 1) + + weight = self.method.get_weight(256, 128, torch.float16) self.assertEqual(weight['weight'].dtype, torch.int8) - self.assertEqual(weight['weight'].shape, (20, 10)) def test_get_pertensor_param(self): - params = self.method.get_pertensor_param(torch.bfloat16) - self.assertEqual(params['input_scale'].dtype, torch.bfloat16) - self.assertEqual(params['input_offset'].dtype, torch.int8) - self.assertEqual(params['input_scale'].shape, (1, )) - self.assertEqual(params['input_offset'].shape, (1, )) + dtypes = [torch.bfloat16, torch.float16, torch.float32] + for dtype in dtypes: + params = self.method.get_pertensor_param(dtype) + self.assertEqual(params['input_scale'].dtype, dtype) + self.assertEqual(params['input_offset'].dtype, torch.int8) + self.assertEqual(params['input_scale'].shape, (1,)) + self.assertEqual(params['input_offset'].shape, (1,)) def test_get_perchannel_param(self): - params = self.method.get_perchannel_param(10, torch.bfloat16) - - self.assertEqual(params['quant_bias'].dtype, torch.int32) - self.assertEqual(params['deq_scale'].dtype, torch.float32) - self.assertEqual(params['weight_scale'].dtype, torch.bfloat16) - self.assertEqual(params['weight_offset'].dtype, torch.bfloat16) - self.assertEqual(params['quant_bias'].shape, (10, )) - self.assertEqual(params['deq_scale'].shape, (10, )) - self.assertEqual(params['weight_scale'].shape, (10, 1)) - self.assertEqual(params['weight_offset'].shape, (10, 1)) + for output_size, dtype in [(128, torch.bfloat16), (256, torch.float16)]: + params = self.method.get_perchannel_param(output_size, dtype) + self.assertEqual(params['quant_bias'].shape, (output_size,)) + self.assertEqual(params['quant_bias'].dtype, torch.int32) + self.assertEqual(params['weight_scale'].shape, (output_size, 1)) + self.assertEqual(params['weight_scale'].dtype, dtype) + self.assertEqual(params['weight_offset'].shape, (output_size, 1)) + self.assertEqual(params['weight_offset'].dtype, dtype) + self.assertEqual(params['deq_scale'].shape, (output_size, )) + if dtype == torch.bfloat16: + self.assertEqual(params['deq_scale'].dtype, torch.float32) + elif dtype == torch.float16: + self.assertEqual(params['deq_scale'].dtype, torch.int64) @patch( "vllm_ascend.quantization.methods.w8a8_static.get_weight_prefetch_method" @@ -49,6 +61,8 @@ def test_apply_with_x_not_int8(self, mock_npu_quant_matmul, mock_quantize, layer.aclnn_input_offset = 0.2 layer.weight = torch.randn(128, 256) layer.deq_scale = 0.3 + quant_bias = torch.zeros(256) + layer.quant_bias = quant_bias mock_get_weight_prefetch_method.return_value = MagicMock() @@ -64,37 +78,21 @@ def test_apply_with_x_not_int8(self, mock_npu_quant_matmul, mock_quantize, output = self.method.apply(layer, x, bias) - expected_y_output += bias - self.assertTrue(torch.equal(output, expected_y_output)) - - @patch("torch_npu.npu_quant_matmul") - def test_apply_with_x_is_int8(self, mock_npu_quant_matmul): - layer = MagicMock() - layer.aclnn_input_scale = 0.1 - layer.aclnn_input_offset = 0.2 - layer.weight = torch.randn(128, 256) - layer.deq_scale = 0.3 - - x = torch.randint(-128, 127, (32, 128), dtype=torch.int8) - bias = torch.randn(256) - - expected_y_output = torch.randn(32, 256) - mock_npu_quant_matmul.return_value = expected_y_output - - output = self.method.apply(layer, x, bias) - expected_y_output += bias self.assertTrue(torch.equal(output, expected_y_output)) + mock_quantize.assert_called_once() + mock_npu_quant_matmul.assert_called_once() + call_kwargs = mock_npu_quant_matmul.call_args.kwargs + self.assertTrue(torch.equal(call_kwargs['bias'], quant_bias)) - @patch('vllm_ascend.utils.get_ascend_device_type', - return_value=AscendDeviceType._310P) + @patch("torch.ops.vllm.quantize") @patch("torch_npu.npu_quant_matmul") - def test_apply_with_x_is_310p(self, mock_npu_quant_matmul, - mock_soc_version): + def test_apply_with_x_is_int8(self, mock_npu_quant_matmul, mock_quantize): layer = MagicMock() layer.aclnn_input_scale = 0.1 layer.aclnn_input_offset = 0.2 layer.weight = torch.randn(128, 256) layer.deq_scale = 0.3 + layer.ascend_quant_method = COMPRESSED_TENSORS_METHOD x = torch.randint(-128, 127, (32, 128), dtype=torch.int8) bias = torch.randn(256) @@ -103,8 +101,11 @@ def test_apply_with_x_is_310p(self, mock_npu_quant_matmul, mock_npu_quant_matmul.return_value = expected_y_output output = self.method.apply(layer, x, bias) - expected_y_output += bias self.assertTrue(torch.equal(output, expected_y_output)) + mock_quantize.assert_not_called() + mock_npu_quant_matmul.assert_called_once() + call_kwargs = mock_npu_quant_matmul.call_args.kwargs + self.assertTrue(torch.equal(call_kwargs['bias'], bias)) @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "0"}) @patch('torch_npu.npu_format_cast') @@ -112,16 +113,15 @@ def test_process_weights_after_loading_with_nz0(self, mock_npu_format_cast): layer = MagicMock() - layer.weight.data = torch.randint(-127, - 128, (128, 256), + layer.weight.data = torch.randint(-128, + 127, (128, 256), dtype=torch.int8) layer.input_scale.data = torch.tensor([0.1]) layer.input_offset.data = torch.tensor([0]) - layer.deq_scale = torch.tensor([0.5]) layer.weight_scale.data = torch.randn(128, 1) layer.weight_offset.data = torch.randn(128, 1) - mock_npu_format_cast.return_value = MagicMock + mock_npu_format_cast.side_effect = identity self.method.process_weights_after_loading(layer) expected_offset = torch.tensor([0]).repeat(256).to(torch.int8) @@ -129,11 +129,11 @@ def test_process_weights_after_loading_with_nz0(self, torch.equal(layer.aclnn_input_offset.data, expected_offset)) self.assertFalse(layer.aclnn_input_offset.requires_grad) - self.assertFalse(layer.deq_scale.requires_grad) - + self.assertEqual(layer.weight.data.shape, (256, 128)) self.assertEqual(layer.weight_scale.data.shape, (128, )) self.assertEqual(layer.weight_offset.data.shape, (128, )) mock_npu_format_cast.assert_not_called() + self.assertNotIn("deq_scale", dict(layer.named_parameters())) @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "1"}) @patch('torch_npu.npu_format_cast') @@ -141,16 +141,15 @@ def test_process_weights_after_loading_with_nz1(self, mock_npu_format_cast): layer = MagicMock() - layer.weight.data = torch.randint(-127, - 128, (128, 256), + layer.weight.data = torch.randint(-128, + 127, (128, 256), dtype=torch.int8) layer.input_scale.data = torch.tensor([0.1]) layer.input_offset.data = torch.tensor([0]) - layer.deq_scale = torch.tensor([0.5]) layer.weight_scale.data = torch.randn(128, 1) layer.weight_offset.data = torch.randn(128, 1) - mock_npu_format_cast.return_value = MagicMock + mock_npu_format_cast.side_effect = identity self.method.process_weights_after_loading(layer) expected_offset = torch.tensor([0]).repeat(256).to(torch.int8) @@ -158,28 +157,28 @@ def test_process_weights_after_loading_with_nz1(self, torch.equal(layer.aclnn_input_offset.data, expected_offset)) self.assertFalse(layer.aclnn_input_offset.requires_grad) - self.assertFalse(layer.deq_scale.requires_grad) - + self.assertEqual(layer.weight.data.shape, (256, 128)) self.assertEqual(layer.weight_scale.data.shape, (128, )) self.assertEqual(layer.weight_offset.data.shape, (128, )) mock_npu_format_cast.assert_called_once() + self.assertTrue(isinstance(layer.deq_scale, MagicMock)) @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "2"}) @patch('torch_npu.npu_format_cast') - def test_process_weights_after_loading_with_nz2(self, - mock_npu_format_cast): + def test_process_weights_after_loading_with_nz2_and_compressed_tensors(self, + mock_npu_format_cast): layer = MagicMock() - layer.weight.data = torch.randint(-127, - 128, (128, 256), + layer.weight.data = torch.randint(-128, + 127, (128, 256), dtype=torch.int8) layer.input_scale.data = torch.tensor([0.1]) layer.input_offset.data = torch.tensor([0]) - layer.deq_scale = torch.tensor([0.5]) layer.weight_scale.data = torch.randn(128, 1) layer.weight_offset.data = torch.randn(128, 1) + layer.ascend_quant_method = COMPRESSED_TENSORS_METHOD - mock_npu_format_cast.return_value = MagicMock + mock_npu_format_cast.side_effect = identity self.method.process_weights_after_loading(layer) expected_offset = torch.tensor([0]).repeat(256).to(torch.int8) @@ -187,8 +186,9 @@ def test_process_weights_after_loading_with_nz2(self, torch.equal(layer.aclnn_input_offset.data, expected_offset)) self.assertFalse(layer.aclnn_input_offset.requires_grad) - self.assertFalse(layer.deq_scale.requires_grad) - + self.assertEqual(layer.weight.data.shape, (256, 128)) self.assertEqual(layer.weight_scale.data.shape, (128, )) self.assertEqual(layer.weight_offset.data.shape, (128, )) mock_npu_format_cast.assert_called_once() + self.assertIn("deq_scale", dict(layer.named_parameters())) + self.assertFalse(isinstance(layer.deq_scale, MagicMock)) diff --git a/tests/ut/quantization/test_compressed_tensors_config.py b/tests/ut/quantization/test_compressed_tensors_config.py new file mode 100644 index 000000000000..28db27fd1bb6 --- /dev/null +++ b/tests/ut/quantization/test_compressed_tensors_config.py @@ -0,0 +1,166 @@ +from unittest.mock import MagicMock, patch + +import torch + +from tests.ut.base import TestBase +from tests.ut.quantization.conftest_quantization import COMPRESSED_TENSORS_W8A8_CONFIG +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.linear import RowParallelLinear, UnquantizedLinearMethod +from vllm_ascend.ops.fused_moe.fused_moe import AscendUnquantizedFusedMoEMethod +from vllm_ascend.quantization.compressed_tensors_config import AscendCompressedTensorsConfig +from vllm_ascend.quantization.method_adapters import AscendLinearMethod, AscendFusedMoEMethod +from vllm_ascend.quantization.methods import AscendW8A8DynamicLinearMethod, AscendW8A8DynamicFusedMoEMethod +from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD + + +class TestAscendCompressedTensorsConfigBasic(TestBase): + + def test_get_name(self): + config = AscendCompressedTensorsConfig.from_config(COMPRESSED_TENSORS_W8A8_CONFIG) + self.assertEqual(config.get_name(), "compressed-tensors") + + def test_get_supported_act_dtypes(self): + dtypes = AscendCompressedTensorsConfig.get_supported_act_dtypes() + self.assertIn(torch.int8, dtypes) + self.assertIn(torch.float16, dtypes) + self.assertIn(torch.bfloat16, dtypes) + self.assertEqual(len(dtypes), 3) + + def test_get_min_capability_raises(self): + with self.assertRaises(NotImplementedError): + AscendCompressedTensorsConfig.get_min_capability() + + def test_get_config_filenames(self): + filenames = AscendCompressedTensorsConfig.get_config_filenames() + self.assertEqual(filenames, []) + + def test_init(self): + config = AscendCompressedTensorsConfig.from_config(COMPRESSED_TENSORS_W8A8_CONFIG) + self.assertEqual(config.ignore, ["lm_head"]) + self.assertEqual(config.quant_format, "int-quantized") + self.assertEqual(list(config.target_scheme_map.keys()), ["Linear"]) + self.assertEqual(config.target_scheme_map["Linear"]["format"], "int-quantized") + self.assertIsNotNone(config.target_scheme_map["Linear"]["input_activations"]) + + def test_apply_vllm_mapper(self): + hf_to_vllm_mapper = MagicMock() + config = AscendCompressedTensorsConfig( + target_scheme_map={"Linear": {}}, + ignore=["lm_head"], + quant_format="", + ) + config.apply_vllm_mapper(hf_to_vllm_mapper) + hf_to_vllm_mapper.apply_dict.assert_called_once() + hf_to_vllm_mapper.apply_list.assert_called_once() + + +class TestAscendCompressedTensorsQuanType(TestBase): + + def setUp(self): + self.config = AscendCompressedTensorsConfig( + target_scheme_map={"Linear": {}}, + ignore=["lm_head"], + quant_format="", + config={}, + ) + + def _make_weight_quant(self, num_bits=8, strategy="channel", dynamic=False, symmetric=True, group_size=None): + mock = MagicMock() + mock.num_bits = num_bits + mock.strategy = strategy + mock.dynamic = dynamic + mock.symmetric = symmetric + mock.group_size = group_size + return mock + + def _make_input_quant(self, num_bits=8, strategy="tensor", dynamic=False, symmetric=True): + mock = MagicMock() + mock.num_bits = num_bits + mock.strategy = strategy + mock.dynamic = dynamic + mock.symmetric = symmetric + return mock + + def test_detect_w8a8_static(self): + weight = self._make_weight_quant(num_bits=8, strategy="channel", dynamic=False, symmetric=True) + input_q = self._make_input_quant(num_bits=8, strategy="tensor", dynamic=False, symmetric=True) + result = self.config._detect_quant_type(weight, input_q, "int-quantized") + self.assertEqual(result, "W8A8") + + def test_detect_w8a8_dynamic(self): + weight = self._make_weight_quant(num_bits=8, strategy="channel", dynamic=False, symmetric=True) + input_q = self._make_input_quant(num_bits=8, strategy="token", dynamic=True, symmetric=True) + result = self.config._detect_quant_type(weight, input_q, "int-quantized") + self.assertEqual(result, "W8A8_DYNAMIC") + + def test_detect_w4a8_dynamic(self): + weight = self._make_weight_quant(num_bits=4, strategy="channel", dynamic=False, symmetric=True) + input_q = self._make_input_quant(num_bits=8, strategy="token", dynamic=True, symmetric=True) + result = self.config._detect_quant_type(weight, input_q, "int-quantized") + self.assertEqual(result, "W4A8_DYNAMIC") + + def test_detect_w4a16(self): + from compressed_tensors.quantization import QuantizationType + weight = MagicMock() + weight.num_bits = 4 + weight.strategy = "group" + weight.dynamic = False + weight.type = QuantizationType.INT + result = self.config._detect_quant_type(weight, None, None) + self.assertEqual(result, "W4A16") + + def test_detect_unsupported_raises(self): + weight = self._make_weight_quant(num_bits=2, strategy="channel", dynamic=False, symmetric=True) + input_q = self._make_input_quant(num_bits=2, strategy="tensor", dynamic=False, symmetric=True) + with self.assertRaises(NotImplementedError): + self.config._detect_quant_type(weight, input_q, "int_quantized") + + +class TestAscendCompressedTensorsConfigGetQuantMethod(TestBase): + + def setUp(self): + self.config = AscendCompressedTensorsConfig.from_config(COMPRESSED_TENSORS_W8A8_CONFIG) + + def test_get_linear_quant_method(self): + layer = MagicMock(spec=RowParallelLinear) + result = self.config.get_quant_method(layer, "model.layers.0.self_attn.q_proj") + self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) + self.assertTrue(isinstance(result, AscendLinearMethod)) + self.assertTrue(isinstance(layer.scheme, AscendW8A8DynamicLinearMethod)) + + def test_get_linear_unquantized_method(self): + layer = MagicMock(spec=RowParallelLinear) + result = self.config.get_quant_method(layer, "lm_head") + self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) + self.assertTrue(isinstance(result, UnquantizedLinearMethod)) + + from vllm_ascend.quantization.methods import AscendW8A8DynamicLinearMethod, AscendW8A8DynamicFusedMoEMethod + @patch("vllm_ascend.quantization.methods.AscendW8A8DynamicFusedMoEMethod.__init__") + def test_get_moe_quant_method(self, mock_method): + mock_method.return_value = None + mock_method.return_value = MagicMock(spec=AscendW8A8DynamicFusedMoEMethod) + layer = MagicMock(spec=FusedMoE) + layer.moe_config = {} + result = self.config.get_quant_method(layer, "model.layers.0.mlp.experts") + self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) + self.assertTrue(isinstance(result, AscendFusedMoEMethod)) + self.assertTrue(isinstance(layer.scheme, AscendW8A8DynamicFusedMoEMethod)) + + @patch("vllm_ascend.ops.fused_moe.fused_moe.AscendUnquantizedFusedMoEMethod.__init__") + @patch("vllm_ascend.quantization.compressed_tensors_config.should_ignore_layer") + def test_get_moe_unquantized_method(self, mock_ignore_layer, mock_method): + mock_method.return_value = None + mock_ignore_layer.return_value = True + layer = MagicMock(spec=FusedMoE) + layer.moe_config = {} + result = self.config.get_quant_method(layer, "model.layers.0.mlp.experts") + self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) + self.assertTrue(isinstance(result, AscendUnquantizedFusedMoEMethod)) + + + def test_no_quant_method(self): + layer = MagicMock(spec=Attention) + result = self.config.get_quant_method(layer, "attn") + self.assertIsNone(result) + diff --git a/tests/ut/quantization/test_method_adapters.py b/tests/ut/quantization/test_method_adapters.py new file mode 100644 index 000000000000..05a6d90604e7 --- /dev/null +++ b/tests/ut/quantization/test_method_adapters.py @@ -0,0 +1,232 @@ +from unittest.mock import MagicMock, patch + +import torch + +from tests.ut.base import TestBase +from vllm.model_executor.layers.fused_moe import FusedMoeWeightScaleSupported +from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm_ascend.quantization.method_adapters import ( + AscendLinearMethod, + AscendKVCacheMethod, + AscendFusedMoEMethod, AscendEmbeddingMethod, +) +from vllm_ascend.quantization.methods.base import AscendLinearScheme, AscendAttentionScheme, AscendMoEScheme + + +class TestAscendLinearMethod(TestBase): + + def setUp(self): + self.mock_scheme = MagicMock(spec=AscendLinearScheme) + self.mock_scheme.get_weight.return_value = { + "weight": torch.empty(128, 256, dtype=torch.int8), + "_packed_dim": 0, + "_packed_factor": 0.1, + } + self.mock_scheme.get_pertensor_param.return_value = { + "weight_scale_pertensor": torch.empty(1, 1, dtype=torch.int8), + } + self.mock_scheme.get_perchannel_param.return_value = { + "weight_scale_perchannel": torch.empty(128, 1, dtype=torch.int8), + } + self.mock_scheme.get_pergroup_param.return_value = { + "weight_scale_second": torch.empty(128, 2, dtype=torch.int8), + "weight_offset_second": torch.empty(128, 2, dtype=torch.int8), + "weight_scale_pergroup": torch.empty(128, 2, dtype=torch.int8), + } + self.method = AscendLinearMethod(self.mock_scheme) + + @patch("vllm_ascend.quantization.method_adapters.PerTensorScaleParameter") + def test_create_weights(self, mock_parameter): + mock_parameter.return_value = torch.nn.Parameter(torch.empty(1, 1, dtype=torch.int8), requires_grad=False) + layer = torch.nn.Module() + weight_loader = MagicMock() + self.method.create_weights( + layer, + input_size_per_partition=256, + output_partition_sizes=[128], + input_size=256, + output_size=128, + params_dtype=torch.bfloat16, + weight_loader=weight_loader, + ) + # Check get_weight method + self.mock_scheme.get_weight.assert_called_once_with(256, 128, torch.bfloat16) + self.assertIn("weight", dict(layer.named_parameters())) + self.assertNotIn("_packed_dim", dict(layer.named_parameters())) + self.assertNotIn("_packed_factor", dict(layer.named_parameters())) + self.assertEqual(layer.weight.input_dim, 1) + self.assertEqual(layer.weight.output_dim, 0) + self.assertEqual(layer.weight.packed_dim, 0) + self.assertEqual(layer.weight.packed_factor, 0.1) + + # Check per tensor param + self.mock_scheme.get_pertensor_param.assert_called_once_with(torch.bfloat16) + self.assertTrue(layer.weight_scale_pertensor.ignore_warning) + self.assertEqual(layer.weight_scale_pertensor.weight_loader, weight_loader) + + # Check per channel param + self.mock_scheme.get_perchannel_param.assert_called_once_with(128, torch.bfloat16) + self.assertEqual(layer.weight_scale_perchannel.output_dim, 0) + self.assertEqual(layer.weight_scale_perchannel.weight_loader, weight_loader) + + # Check per group param + self.mock_scheme.get_pergroup_param.assert_called_once() + self.assertEqual(layer.weight_scale_pergroup.output_dim, 0) + self.assertFalse(hasattr(layer.weight_scale_pergroup, "input_dim")) + self.assertEqual(layer.weight_scale_second.input_dim, 1) + self.assertEqual(layer.weight_offset_second.input_dim, 1) + + def test_process_weights_after_loading_delegates(self): + layer = torch.nn.Module() + self.mock_scheme.process_weights_after_loading.return_value = None + self.method.process_weights_after_loading(layer) + self.mock_scheme.process_weights_after_loading.assert_called_once_with(layer) + + def test_process_weights_after_loading_skips_when_not_present(self): + del self.mock_scheme.process_weights_after_loading + layer = torch.nn.Module() + self.method.process_weights_after_loading(layer) + + def test_get_computed_params(self): + computed = self.method.get_computed_params() + self.assertIn("weight_offset", computed) + self.assertIn("quant_bias", computed) + self.assertIn("deq_scale", computed) + self.assertIn("weight_scale", computed) + self.assertEqual(len(computed), 4) + + def test_apply_delegates_to_scheme(self): + layer = MagicMock(spec=ColumnParallelLinear) + x = torch.randn(4, 256) + self.mock_scheme.apply.return_value = torch.randn(4, 128) + output = self.method.apply(layer, x) + self.mock_scheme.apply.assert_called_once() + self.assertEqual(output.shape, (4, 128)) + + +class TestAscendKVCacheMethod(TestBase): + + def setUp(self): + self.mock_scheme = MagicMock(spec=AscendAttentionScheme) + self.mock_scheme.create_weights.return_value = None + self.mock_scheme.process_weights_after_loading.return_value = None + self.method = AscendKVCacheMethod(self.mock_scheme) + + def test_create_weights_delegates(self): + layer = torch.nn.Module() + self.method.create_weights(layer) + self.mock_scheme.create_weights.assert_called_once_with(layer) + + def test_process_weights_after_loading_delegates(self): + layer = torch.nn.Module() + self.method.process_weights_after_loading(layer) + self.mock_scheme.process_weights_after_loading.assert_called_once_with(layer) + + def test_apply_delegates(self): + layer = torch.nn.Module() + query = torch.randn(4, 8, 64) + key = torch.randn(4, 8, 64) + value = torch.randn(4, 8, 64) + self.mock_scheme.apply.return_value = torch.randn(4, 8, 64) + self.method.apply( + layer, query, key, value, + kv_cache=None, attn_metadata=None, + attn_type=None, scale=1.0, output=None, + ) + self.mock_scheme.apply.assert_called_once() + + +class TestAscendFusedMoEMethod(TestBase): + + def setUp(self): + self.mock_scheme = MagicMock(spec=AscendMoEScheme) + self.mock_scheme.group_size = 0 + self.mock_moe_config = MagicMock() + self.method = AscendFusedMoEMethod(self.mock_scheme, self.mock_moe_config) + + def test_process_weights_after_loading_delegates(self): + layer = torch.nn.Module() + self.mock_scheme.process_weights_after_loading.return_value = None + self.method.process_weights_after_loading(layer) + self.mock_scheme.process_weights_after_loading.assert_called_once_with(layer) + + def test_process_weights_after_loading_skips_when_not_present(self): + del self.mock_scheme.process_weights_after_loading + layer = torch.nn.Module() + self.method.process_weights_after_loading(layer) + + def test_create_weights_registers_parameters(self): + self.mock_scheme.get_weight.return_value = { + "w13_weight": torch.empty(8, 256, 128, dtype=torch.int8), + "w2_weight": torch.empty(8, 128, 256, dtype=torch.int8), + } + self.mock_scheme.get_dynamic_quant_param.return_value = { + "w13_weight_scale_second": torch.empty(8, 256, 1, dtype=torch.bfloat16), + "w2_weight_offset_second": torch.empty(8, 128, 1, dtype=torch.bfloat16), + "w2_scale_bias": torch.empty(8, 128, 1, dtype=torch.bfloat16), + "w13_weight_scale": torch.empty(8, 256, 1, dtype=torch.bfloat16), + "w2_weight_offset": torch.empty(8, 128, 1, dtype=torch.bfloat16), + } + # per channel quantization + layer = torch.nn.Module() + self.method.create_weights( + layer, + num_experts=8, + hidden_size=128, + intermediate_size_per_partition=256, + params_dtype=torch.bfloat16, + ) + self.assertIn("w13_weight", dict(layer.named_parameters())) + self.assertIn("w2_weight", dict(layer.named_parameters())) + + self.assertEqual(layer.w13_weight_scale_second.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + self.assertEqual(layer.w2_weight_offset_second.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + self.assertEqual(layer.w2_scale_bias.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + self.assertEqual(layer.w13_weight_scale.quant_method, FusedMoeWeightScaleSupported.CHANNEL.value) + self.assertEqual(layer.w2_weight_offset.quant_method, FusedMoeWeightScaleSupported.CHANNEL.value) + + # per group quantization + self.mock_scheme.group_size = 128 + layer = torch.nn.Module() + self.method.create_weights( + layer, + num_experts=8, + hidden_size=128, + intermediate_size_per_partition=256, + params_dtype=torch.bfloat16, + ) + self.assertEqual(layer.w13_weight_scale_second.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + self.assertEqual(layer.w2_weight_offset_second.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + self.assertEqual(layer.w2_scale_bias.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + self.assertEqual(layer.w13_weight_scale.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + self.assertEqual(layer.w2_weight_offset.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + + + + def test_apply_method(self): + layer = torch.nn.Module() + x = torch.randn(8, 64) + router_logits = torch.randn(8, 64) + top_k = 3 + renormalize = True + self.mock_scheme.apply.return_value = None + self.method.apply( + layer, x, router_logits, top_k, renormalize + ) + self.mock_scheme.apply.assert_called_once() + + def test_supports_eplb_default_false(self): + self.assertFalse(self.method.supports_eplb) + + def test_supports_eplb_when_scheme_supports(self): + self.mock_scheme.supports_eplb = True + self.assertTrue(self.method.supports_eplb) + + +class TestAscendEmbeddingMethod(TestBase): + + def test_init(self): + layer = MagicMock(spec=AscendLinearScheme) + method = AscendEmbeddingMethod(layer) + self.assertTrue(isinstance(method, AscendLinearMethod)) + self.assertEqual(method.quant_method, layer) diff --git a/tests/ut/quantization/test_modelslim_config.py b/tests/ut/quantization/test_modelslim_config.py index 8970596f57cb..1161b7f6cd9f 100644 --- a/tests/ut/quantization/test_modelslim_config.py +++ b/tests/ut/quantization/test_modelslim_config.py @@ -3,6 +3,8 @@ import tempfile from unittest.mock import MagicMock, patch +import torch + from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.linear import LinearBase @@ -12,6 +14,10 @@ from vllm_ascend.quantization.modelslim_config import ( MODELSLIM_CONFIG_FILENAME, AscendModelSlimConfig, + create_scheme_for_layer, + get_linear_quant_type, + get_packed_modules_mapping, + get_quant_type_for_layer, ) from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD @@ -267,3 +273,407 @@ def test_apply_extra_quant_adaptations_weight_packed(self): self.assertIn("model.layers.0.weight", config.quant_description) self.assertEqual(config.quant_description["model.layers.0.weight"], "INT8") + + +class TestGetPackedModulesMapping(TestBase): + + def test_known_model_type(self): + result = get_packed_modules_mapping("deepseek_v3") + self.assertIsInstance(result, dict) + self.assertIn("gate_up_proj", result) + self.assertIn("experts", result) + + def test_unknown_model_type(self): + result = get_packed_modules_mapping("unknown_model") + self.assertEqual(result, {}) + + +class TestGetLinearQuantType(TestBase): + + def test_simple_non_fused_layer(self): + quant_description = {"layer1.weight": "W8A8_DYNAMIC"} + packed_modules_mapping = {} + result = get_linear_quant_type(quant_description, "layer1", packed_modules_mapping) + self.assertEqual(result, "W8A8_DYNAMIC") + + def test_fused_layer_consistent_quant_types(self): + quant_description = { + "shard1.weight": "W8A8_DYNAMIC", + "shard2.weight": "W8A8_DYNAMIC", + } + packed_modules_mapping = {"fused_layer": ["shard1", "shard2"]} + result = get_linear_quant_type(quant_description, "fused_layer", packed_modules_mapping) + self.assertEqual(result, "W8A8_DYNAMIC") + + def test_fused_layer_inconsistent_quant_types(self): + quant_description = { + "shard1.weight": "W8A8_DYNAMIC", + "shard2.weight": "FLOAT", + } + packed_modules_mapping = {"fused_layer": ["shard1", "shard2"]} + with self.assertRaises(ValueError) as ctx: + get_linear_quant_type(quant_description, "fused_layer", packed_modules_mapping) + self.assertIn("Not all shards", str(ctx.exception)) + + +class TestGetQuantTypeForLayer(TestBase): + + def test_attention_layer_with_fa_quant_type(self): + quant_description = { + "fa_quant_type": "C8", + "layer1.weight": "W8A8_DYNAMIC", + } + result = get_quant_type_for_layer(quant_description, "layer1", "attention") + self.assertEqual(result, "C8") + + def test_attention_layer_with_indexer_quant_type(self): + quant_description = { + "indexer_quant_type": "INT8", + "layer1.weight": "W8A8_DYNAMIC", + } + result = get_quant_type_for_layer(quant_description, "layer1", "attention") + self.assertEqual(result, "INT8") + + def test_attention_layer_with_layer_indexer_quant_type(self): + quant_description = { + "layer1.indexer.quant_type": "INT8", + "layer1.weight": "W8A8_DYNAMIC", + } + result = get_quant_type_for_layer(quant_description, "layer1", "attention") + self.assertEqual(result, "INT8") + + def test_linear_layer(self): + quant_description = { + "layer1.weight": "W8A8_DYNAMIC", + } + result = get_quant_type_for_layer(quant_description, "layer1", "linear") + self.assertEqual(result, "W8A8_DYNAMIC") + + def test_moe_layer(self): + quant_description = { + "moe_layer.weight": "W8A8_DYNAMIC", + } + result = get_quant_type_for_layer(quant_description, "moe_layer", "moe") + self.assertEqual(result, "W8A8_DYNAMIC") + + def test_with_none_packed_modules_mapping(self): + quant_description = { + "layer1.weight": "W8A8_DYNAMIC", + } + result = get_quant_type_for_layer(quant_description, "layer1", "linear", None) + self.assertEqual(result, "W8A8_DYNAMIC") + + +class TestCreateSchemeForLayer(TestBase): + + @patch("vllm_ascend.quantization.modelslim_config.get_scheme_class") + def test_successful_scheme_creation(self, mock_get_scheme_class): + mock_scheme_class = MagicMock() + mock_scheme_instance = MagicMock() + mock_scheme_class.return_value = mock_scheme_instance + mock_get_scheme_class.return_value = mock_scheme_class + + quant_description = {"layer1.weight": "W8A8_DYNAMIC"} + result = create_scheme_for_layer(quant_description, "layer1", "linear") + + self.assertIs(result, mock_scheme_instance) + mock_get_scheme_class.assert_called_once_with("W8A8_DYNAMIC", "linear") + + def test_quant_type_none_raises_error(self): + quant_description = {} + with self.assertRaises(KeyError) as ctx: + create_scheme_for_layer(quant_description, "layer1", "linear") + + @patch("vllm_ascend.quantization.modelslim_config.get_scheme_class") + def test_unsupported_quant_type_raises_error(self, mock_get_scheme_class): + mock_get_scheme_class.return_value = None + + quant_description = {"layer1.weight": "UNSUPPORTED_TYPE"} + with self.assertRaises(NotImplementedError) as ctx: + create_scheme_for_layer(quant_description, "layer1", "linear") + self.assertIn("doesn't support UNSUPPORTED_TYPE", str(ctx.exception)) + + +class TestApplyVllmMapper(TestBase): + + def test_apply_mapper_with_populated_quant_description(self): + config = AscendModelSlimConfig({"old_key.weight": "INT8"}) + mock_mapper = MagicMock() + mock_mapper.apply_dict.return_value = {"new_key.weight": "INT8"} + + config.apply_vllm_mapper(mock_mapper) + + self.assertEqual(config.quant_description, {"new_key.weight": "INT8"}) + mock_mapper.apply_dict.assert_called_once_with({"old_key.weight": "INT8"}) + + def test_idempotency_when_mapper_already_applied(self): + config = AscendModelSlimConfig({"key.weight": "INT8"}) + mock_mapper = MagicMock() + mock_mapper.apply_dict.return_value = {"key.weight": "INT8"} + + config.apply_vllm_mapper(mock_mapper) + first_result = config.quant_description + + config.apply_vllm_mapper(mock_mapper) + second_result = config.quant_description + + self.assertEqual(first_result, second_result) + mock_mapper.apply_dict.assert_called_once() + + def test_empty_quant_description(self): + config = AscendModelSlimConfig({}) + mock_mapper = MagicMock() + + config.apply_vllm_mapper(mock_mapper) + + mock_mapper.apply_dict.assert_not_called() + + +class TestGetCacheScale(TestBase): + + def test_c8_kv_cache_type_k_proj_scale(self): + config = AscendModelSlimConfig({"kv_cache_type": "C8"}) + result = config.get_cache_scale("model.layers.0.k_proj.kv_cache_scale") + self.assertEqual(result, "model.layers.0.attn.k_cache_scale") + + def test_c8_kv_cache_type_k_proj_offset(self): + config = AscendModelSlimConfig({"kv_cache_type": "C8"}) + result = config.get_cache_scale("model.layers.0.k_proj.kv_cache_offset") + self.assertEqual(result, "model.layers.0.attn.k_cache_offset") + + def test_c8_kv_cache_type_v_proj_scale(self): + config = AscendModelSlimConfig({"kv_cache_type": "C8"}) + result = config.get_cache_scale("model.layers.0.v_proj.kv_cache_scale") + self.assertEqual(result, "model.layers.0.attn.v_cache_scale") + + def test_c8_kv_cache_type_v_proj_offset(self): + config = AscendModelSlimConfig({"kv_cache_type": "C8"}) + result = config.get_cache_scale("model.layers.0.v_proj.kv_cache_offset") + self.assertEqual(result, "model.layers.0.attn.v_cache_offset") + + def test_non_c8_kv_cache_type(self): + config = AscendModelSlimConfig({"kv_cache_type": "FLOAT"}) + result = config.get_cache_scale("model.layers.0.k_proj.kv_cache_scale") + self.assertIsNone(result) + + def test_no_matching_suffix(self): + config = AscendModelSlimConfig({"kv_cache_type": "C8"}) + result = config.get_cache_scale("model.layers.0.other_key") + self.assertIsNone(result) + + +class TestQuantPrefixMapper(TestBase): + + def test_sets_model_type_and_returns_prefix(self): + config = AscendModelSlimConfig({}) + result = config.quant_prefix_mapper("deepseek_v3", "model.layers.0") + self.assertEqual(result, "model.layers.0") + self.assertEqual(config.model_type, "deepseek_v3") + + +class TestIsFaQuantLayer(TestBase): + + def test_enable_fa_quant_true_matching_layer(self): + config = AscendModelSlimConfig({ + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + }) + result = config.is_fa_quant_layer("layers.1.attn") + self.assertTrue(result) + + def test_enable_fa_quant_false(self): + config = AscendModelSlimConfig({}) + result = config.is_fa_quant_layer("layers.1.attn") + self.assertFalse(result) + + def test_non_matching_layer_index(self): + config = AscendModelSlimConfig({ + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + }) + result = config.is_fa_quant_layer("layers.2.attn") + self.assertFalse(result) + + def test_no_layer_index(self): + config = AscendModelSlimConfig({ + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + }) + result = config.is_fa_quant_layer("model.attn") + self.assertFalse(result) + + +class TestEnablingFaQuant(TestBase): + + def test_decode_instance_with_fa_quant_layer(self): + config = AscendModelSlimConfig({ + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + }) + mock_vllm_config = MagicMock() + mock_vllm_config.kv_transfer_config.is_kv_consumer = True + mock_vllm_config.kv_transfer_config.is_kv_producer = False + + result = config.enabling_fa_quant(mock_vllm_config, "layers.1.attn") + self.assertTrue(result) + + def test_non_decode_instance(self): + config = AscendModelSlimConfig({ + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + }) + mock_vllm_config = MagicMock() + mock_vllm_config.kv_transfer_config.is_kv_consumer = False + mock_vllm_config.kv_transfer_config.is_kv_producer = True + + result = config.enabling_fa_quant(mock_vllm_config, "layers.1.attn") + self.assertFalse(result) + + def test_is_fa_quant_layer_false(self): + config = AscendModelSlimConfig({ + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + }) + mock_vllm_config = MagicMock() + mock_vllm_config.kv_transfer_config.is_kv_consumer = True + mock_vllm_config.kv_transfer_config.is_kv_producer = False + + result = config.enabling_fa_quant(mock_vllm_config, "layers.2.attn") + self.assertFalse(result) + + +class TestIsIndexerQuantLayer(TestBase): + + def test_enable_indexer_quant_true_matching_layer(self): + config = AscendModelSlimConfig({ + "indexer_quant_type": "INT8", + "layers.1.indexer.quant_type": "INT8", + }) + result = config.is_indexer_quant_layer("layers.1.attn") + self.assertTrue(result) + + def test_enable_indexer_quant_false(self): + config = AscendModelSlimConfig({}) + result = config.is_indexer_quant_layer("layers.1.attn") + self.assertFalse(result) + + def test_non_matching_layer_index(self): + config = AscendModelSlimConfig({ + "indexer_quant_type": "INT8", + "layers.1.indexer.quant_type": "INT8", + }) + result = config.is_indexer_quant_layer("layers.2.attn") + self.assertFalse(result) + + def test_no_layer_index(self): + config = AscendModelSlimConfig({ + "indexer_quant_type": "INT8", + "layers.1.indexer.quant_type": "INT8", + }) + result = config.is_indexer_quant_layer("model.attn") + self.assertFalse(result) + + +class TestGetKvQuantDtype(TestBase): + + def test_enable_fa_quant_true_with_mla(self): + config = AscendModelSlimConfig({ + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + }) + mock_model_config = MagicMock() + mock_model_config.dtype = torch.float16 + mock_model_config.use_mla = True + + k_dtype, v_dtype = config.get_kv_quant_dtype("layers.1.attn", torch.float16, mock_model_config) + self.assertEqual(k_dtype, torch.int8) + self.assertEqual(v_dtype, torch.float16) + + def test_enable_fa_quant_true_without_mla(self): + config = AscendModelSlimConfig({ + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + }) + mock_model_config = MagicMock() + mock_model_config.dtype = torch.float16 + mock_model_config.use_mla = False + + k_dtype, v_dtype = config.get_kv_quant_dtype("layers.1.attn", torch.float16, mock_model_config) + self.assertEqual(k_dtype, torch.int8) + self.assertEqual(v_dtype, torch.int8) + + def test_enable_fa_quant_false(self): + config = AscendModelSlimConfig({}) + mock_model_config = MagicMock() + mock_model_config.dtype = torch.float16 + mock_model_config.use_mla = True + + k_dtype, v_dtype = config.get_kv_quant_dtype("layers.1.attn", torch.float16, mock_model_config) + self.assertEqual(k_dtype, torch.float16) + self.assertEqual(v_dtype, torch.float16) + + +class TestGetKvQuantSplitFactor(TestBase): + + @patch("vllm_ascend.quantization.modelslim_config.calc_split_factor") + def test_enable_fa_quant_true(self, mock_calc_split_factor): + mock_calc_split_factor.return_value = 2.0 + config = AscendModelSlimConfig({ + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + }) + kv_head_dim_list = [64, 64] + + result = config.get_kv_quant_split_factor("layers.1.attn", kv_head_dim_list) + self.assertEqual(result, 2.0) + mock_calc_split_factor.assert_called_once_with([64, 128]) + + @patch("vllm_ascend.quantization.modelslim_config.calc_split_factor") + def test_enable_fa_quant_false(self, mock_calc_split_factor): + mock_calc_split_factor.return_value = 1.0 + config = AscendModelSlimConfig({}) + kv_head_dim_list = [64, 64] + + result = config.get_kv_quant_split_factor("layers.1.attn", kv_head_dim_list) + self.assertEqual(result, 1.0) + mock_calc_split_factor.assert_called_once_with([64, 64]) + + +class TestAddKvcacheQuantMetadata(TestBase): + + def test_with_fa_quant_type(self): + config = AscendModelSlimConfig({ + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + "layers.2.fa_k.scale": "C8", + }) + config._add_kvcache_quant_metadata() + + self.assertTrue(config.enable_fa_quant) + self.assertIn(1, config.kvcache_quant_layers) + self.assertIn(2, config.kvcache_quant_layers) + self.assertFalse(config.enable_indexer_quant) + self.assertEqual(config.indexer_quant_layers, []) + + def test_with_indexer_quant_type(self): + config = AscendModelSlimConfig({ + "indexer_quant_type": "INT8", + "layers.1.indexer.quant_type": "INT8", + "layers.3.indexer.quant_type": "INT8", + }) + config._add_kvcache_quant_metadata() + + self.assertFalse(config.enable_fa_quant) + self.assertEqual(config.kvcache_quant_layers, []) + self.assertTrue(config.enable_indexer_quant) + self.assertIn(1, config.indexer_quant_layers) + self.assertIn(3, config.indexer_quant_layers) + + def test_with_neither_quant_type(self): + config = AscendModelSlimConfig({}) + config._add_kvcache_quant_metadata() + + self.assertFalse(config.enable_fa_quant) + self.assertEqual(config.kvcache_quant_layers, []) + self.assertFalse(config.enable_indexer_quant) + self.assertEqual(config.indexer_quant_layers, []) diff --git a/tests/ut/quantization/test_quant_parser.py b/tests/ut/quantization/test_quant_parser.py new file mode 100644 index 000000000000..75e544c89f11 --- /dev/null +++ b/tests/ut/quantization/test_quant_parser.py @@ -0,0 +1,123 @@ +from unittest.mock import patch +import torch +from tests.ut.base import TestBase +from vllm_ascend.quantization.quant_parser import ( + QuantTypeMapping, + get_rollback_quant_type, + parse_mxfp_quant_params, + parse_quant_moe_down_proj_params, +) + + +class TestQuantTypeMapping(TestBase): + + def test_get_quant_settings_returns_dict(self): + settings = QuantTypeMapping.get_quant_settings() + self.assertIsInstance(settings, dict) + + def test_get_quant_settings_contains_expected_keys(self): + settings = QuantTypeMapping.get_quant_settings() + self.assertIn("W8A8_MXFP8", settings) + self.assertIn("W4A4_MXFP4", settings) + self.assertIn("W4A8_MXFP", settings) + + def test_w8a8_mxfp8_settings(self): + settings = QuantTypeMapping.get_quant_settings()["W8A8_MXFP8"] + self.assertEqual(settings["act_quant_type"], torch.float8_e4m3fn) + self.assertIsNone(settings["weight_quant_type"]) + + def test_w4a4_mxfp4_settings(self): + settings = QuantTypeMapping.get_quant_settings()["W4A4_MXFP4"] + self.assertIn("act_quant_type", settings) + self.assertIn("weight_quant_type", settings) + self.assertIn("scale_dtype", settings) + self.assertIn("per_token_scale_dtype", settings) + + +class TestGetRollbackQuantType(TestBase): + + def test_returns_down_proj_quant_type(self): + config = { + "model.layers.0.mlp.gate_proj": "W8A8_MXFP8", + "model.layers.0.mlp.down_proj": "W4A4_MXFP4", + } + result = get_rollback_quant_type(config) + self.assertEqual(result, "W4A4_MXFP4") + + def test_returns_default_when_no_down_proj(self): + config = {"model.layers.0.mlp.gate_proj": "W4A8_MXFP"} + result = get_rollback_quant_type(config) + self.assertEqual(result, "W8A8_MXFP8") + + def test_returns_down_proj_type_with_multiple_entries(self): + config = { + "model.layers.0.mlp.gate_proj": "W8A8_MXFP8", + "model.layers.0.mlp.up_proj": "W8A8_MXFP8", + "model.layers.0.mlp.down_proj": "W4A8_MXFP", + } + result = get_rollback_quant_type(config) + self.assertEqual(result, "W4A8_MXFP") + + +class TestParseMxfpQuantParams(TestBase): + + def test_default_values(self): + act, weight, scale, per_token, round_mode = parse_mxfp_quant_params() + self.assertEqual(act, torch.float8_e4m3fn) + self.assertEqual(weight, torch.float8_e4m3fn) + self.assertIsNone(scale) + self.assertIsNone(per_token) + self.assertEqual(round_mode, "rint") + + def test_custom_values(self): + act, weight, scale, per_token, round_mode = parse_mxfp_quant_params( + act_quant_type=torch.float16, + weight_quant_type=torch.float8_e4m3fn, + round_mode="round", + ) + self.assertEqual(act, torch.float16) + self.assertEqual(weight, torch.float8_e4m3fn) + self.assertEqual(round_mode, "round") + + def test_scale_type_overrides(self): + act, weight, scale, per_token, round_mode = parse_mxfp_quant_params( + scale_type="float8_e8m0fnu", + per_token_scale_type="float8_e8m0fnu", + ) + self.assertEqual(scale, "float8_e8m0fnu") + self.assertEqual(per_token, "float8_e8m0fnu") + + +class TestParseQuantMoeDownProjParams(TestBase): + + @patch("vllm_ascend.quantization.quant_parser.ensure_mxfp8_scale_dtype_available") + def test_w8a8_mxfp8_uses_rint_round_mode(self, mock_ensure): + mock_ensure.return_value = None + act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params( + "W8A8_MXFP8", "round" + ) + self.assertEqual(round_mode, "rint") + + @patch("vllm_ascend.quantization.quant_parser.ensure_mxfp4_dtype_available") + def test_w4a4_mxfp4_respects_parsed_round_mode(self, mock_ensure): + mock_ensure.return_value = None + act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params( + "W4A4_MXFP4", "round" + ) + self.assertEqual(round_mode, "round") + + @patch("vllm_ascend.quantization.quant_parser.ensure_mxfp4_dtype_available") + def test_w4a4_mxfp4_rint_round_mode(self, mock_ensure): + mock_ensure.return_value = None + act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params( + "W4A4_MXFP4", "rint" + ) + self.assertEqual(round_mode, "rint") + + @patch("vllm_ascend.quantization.quant_parser.ensure_mxfp8_scale_dtype_available") + def test_w4a8_mxfp_uses_rint_round_mode(self, mock_ensure): + mock_ensure.return_value = None + act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params( + "W4A8_MXFP", "round" + ) + self.assertEqual(round_mode, "rint") diff --git a/tests/ut/quantization/test_quant_type.py b/tests/ut/quantization/test_quant_type.py new file mode 100644 index 000000000000..820acd641f25 --- /dev/null +++ b/tests/ut/quantization/test_quant_type.py @@ -0,0 +1,29 @@ +from tests.ut.base import TestBase +from vllm_ascend.quantization.quant_type import QuantType + + +class TestQuantType(TestBase): + + def test_enum_values(self): + self.assertEqual(QuantType.NONE.value, 0) + self.assertEqual(QuantType.W8A8.value, 1) + self.assertEqual(QuantType.W4A8.value, 2) + self.assertEqual(QuantType.MXFP8.value, 3) + self.assertEqual(QuantType.W4A16.value, 4) + self.assertEqual(QuantType.MXFP4.value, 5) + + def test_enum_members_count(self): + members = list(QuantType) + self.assertEqual(len(members), 6) + + def test_enum_identity(self): + self.assertIs(QuantType.W8A8, QuantType.W8A8) + self.assertIsNot(QuantType.W8A8, QuantType.W4A8) + + def test_enum_by_value(self): + self.assertEqual(QuantType(1), QuantType.W8A8) + self.assertEqual(QuantType(3), QuantType.MXFP8) + + def test_enum_by_name(self): + self.assertEqual(QuantType["NONE"], QuantType.NONE) + self.assertEqual(QuantType["W4A16"], QuantType.W4A16) diff --git a/tests/ut/quantization/test_quant_utils.py b/tests/ut/quantization/test_utils.py similarity index 86% rename from tests/ut/quantization/test_quant_utils.py rename to tests/ut/quantization/test_utils.py index 551cf4fff939..45f23a9a700d 100644 --- a/tests/ut/quantization/test_quant_utils.py +++ b/tests/ut/quantization/test_utils.py @@ -1,13 +1,15 @@ import json import os import tempfile -from pathlib import Path from unittest.mock import MagicMock, patch from tests.ut.base import TestBase -from vllm_ascend.quantization.modelslim_config import MODELSLIM_CONFIG_FILENAME +from tests.ut.quantization.conftest_quantization import FAKQUANT_CONFIG, W8A8_CONFIG +from vllm_ascend.quantization import AscendCompressedTensorsConfig +from vllm_ascend.quantization.modelslim_config import MODELSLIM_CONFIG_FILENAME, AscendModelSlimConfig from vllm_ascend.quantization.utils import ( detect_quantization_method, + enable_fa_quant, maybe_auto_detect_quantization, ) from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD, COMPRESSED_TENSORS_METHOD @@ -188,3 +190,33 @@ def test_passes_revision_to_detect(self, mock_detect, mock_get_quant): model_path="org/model-name", revision="v1.0", quantization=None) maybe_auto_detect_quantization(vllm_config) mock_detect.assert_called_once_with("org/model-name", revision="v1.0") + + +class TestEnableFaQuant(TestBase): + + def test_non_quantization_scenarios(self): + vllm_config = MagicMock() + vllm_config.quant_config = None + result = enable_fa_quant(vllm_config) + self.assertFalse(result) + + def test_llmcompressor_quantization_scenario(self): + vllm_config = MagicMock() + vllm_config.quant_config = AscendCompressedTensorsConfig({}, [], "", {}) + result = enable_fa_quant(vllm_config) + self.assertFalse(result) + + def test_not_fa3_quantization_scenario(self): + vllm_config = MagicMock() + vllm_config.quant_config = AscendModelSlimConfig(W8A8_CONFIG) + result = enable_fa_quant(vllm_config) + self.assertFalse(result) + + def test_fa3_quantization_scenario(self): + vllm_config = MagicMock() + vllm_config.quant_config = AscendModelSlimConfig(FAKQUANT_CONFIG) + vllm_config.kv_transfer_config = None + result = enable_fa_quant(vllm_config) + self.assertTrue(result) + result = enable_fa_quant(vllm_config, layer_name="test_layer") + self.assertFalse(result) diff --git a/tests/ut/quantization/test_w8a8_dynamic.py b/tests/ut/quantization/test_w8a8_dynamic.py deleted file mode 100644 index f01898f0d3d1..000000000000 --- a/tests/ut/quantization/test_w8a8_dynamic.py +++ /dev/null @@ -1,229 +0,0 @@ -from unittest.mock import Mock, patch - -import torch - -from tests.ut.base import TestBase -from vllm_ascend.ascend_forward_context import MoECommType -from vllm_ascend.quantization.methods.w8a8_dynamic import AscendW8A8DynamicFusedMoEMethod - - -class TestAscendW8A8FusedMoEMethod(TestBase): - num_experts = 8 - hidden_size = 128 - intermediate_size = 128 - - @patch("torch.distributed.get_rank") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_mc2_group") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ascend_config") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ep_group") - def setUp(self, mock_get_ep_group, mock_get_ascend_config, - mock_get_mc2_group, mock_get_rank): - with patch( - 'vllm_ascend.quantization.methods.w8a8_dynamic.get_current_vllm_config' - ) as mock_get_current_vllm_config: - 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 - mock_ep_group = Mock() - mock_get_ep_group.return_value = mock_ep_group - mock_ascend_config = Mock() - mock_ascend_config.enable_chunked_prefill = False - mock_ascend_config.multistream_overlap_gate = False - mock_ascend_config.eplb_config = Mock(dynamic_eplb=False) - mock_get_ascend_config.return_value = mock_ascend_config - mock_mc2_group = Mock(device_group=0) - mock_get_mc2_group.return_value = mock_mc2_group - mock_rank = Mock() - mock_get_rank.return_value = mock_rank - - self.quant_method = AscendW8A8DynamicFusedMoEMethod() - - def test_get_weight(self): - param_dict = self.quant_method.get_weight(self.num_experts, - self.intermediate_size, - self.hidden_size, - torch.bfloat16) - self.assertEqual(param_dict["w13_weight"].dtype, torch.int8) - self.assertEqual( - param_dict["w13_weight"].shape, - (self.num_experts, 2 * self.intermediate_size, self.hidden_size)) - - def test_get_dynamic_quant_param(self): - param_dict = self.quant_method.get_dynamic_quant_param( - self.num_experts, self.intermediate_size, self.hidden_size, - torch.bfloat16) - self.assertEqual(param_dict["w13_weight_scale"].dtype, torch.bfloat16) - self.assertEqual(param_dict["w13_weight_scale"].shape, - (self.num_experts, 2 * self.intermediate_size, 1)) - - def build_layer(self): - layer = torch.nn.Module() - layer.w13_weight = torch.nn.Parameter(torch.empty( - self.num_experts, - 2 * self.intermediate_size, - self.hidden_size, - dtype=torch.int8), - requires_grad=False) - layer.w2_weight = torch.nn.Parameter(torch.empty( - self.num_experts, - self.hidden_size, - self.intermediate_size, - dtype=torch.int8), - requires_grad=False) - w13_weight_scale = torch.zeros( - (self.num_experts, 2 * self.intermediate_size, 1), - dtype=torch.float32) - layer.w13_weight_scale = torch.nn.Parameter(w13_weight_scale, - requires_grad=False) - w13_weight_offset = torch.zeros( - (self.num_experts, 2 * self.intermediate_size, 1), - dtype=torch.float32) - layer.w13_weight_offset = torch.nn.Parameter(w13_weight_offset, - requires_grad=False) - w2_weight_scale = torch.zeros((self.num_experts, self.hidden_size, 1), - dtype=torch.float32) - layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, - requires_grad=False) - w2_weight_offset = torch.zeros((self.num_experts, self.hidden_size, 1), - dtype=torch.float32) - layer.w2_weight_offset = torch.nn.Parameter(w2_weight_offset, - requires_grad=False) - return layer - - @patch('torch_npu.npu_format_cast') - def test_process_weights_after_loading(self, mock_npu_format_cast): - - def func_by_args(weight, num_format): - return weight - - mock_npu_format_cast.side_effect = func_by_args - new_layer = self.build_layer() - self.quant_method.process_weights_after_loading(new_layer) - mock_npu_format_cast.assert_called() - - @patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts") - def test_apply_uses_explicit_dispatch_and_mlp_args(self, mock_select_experts, mock_extra_ctx): - tokens = 4 - hidden_size = self.hidden_size - layer = torch.nn.Module() - layer.w13_weight = torch.randint( - -8, - 8, - (self.num_experts, 2 * self.intermediate_size, hidden_size), - dtype=torch.int8, - ) - layer.w2_weight = torch.randint( - -8, - 8, - (self.num_experts, hidden_size, self.intermediate_size), - dtype=torch.int8, - ) - layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size, dtype=torch.float32) - layer.w2_weight_scale = torch.ones(self.num_experts, hidden_size, dtype=torch.float32) - - x = torch.randn(tokens, hidden_size, dtype=torch.float32) - router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32) - topk_weights = torch.randn(tokens, 2, dtype=torch.float32) - topk_ids = torch.randint(0, self.num_experts, (tokens, 2), dtype=torch.int64) - mc2_mask = torch.tensor([1, 0, 1, 0], dtype=torch.bool) - pertoken_scale = torch.randn(tokens, dtype=torch.float32) - - mock_select_experts.return_value = (topk_weights, topk_ids) - mock_comm = Mock() - mock_comm.fused_experts.return_value = torch.randn(tokens, hidden_size, dtype=torch.float32) - mock_extra_ctx.moe_comm_method = mock_comm - mock_extra_ctx.moe_comm_type = MoECommType.ALLGATHER - self.quant_method.multistream_overlap_gate = False - self.quant_method.in_dtype = torch.float32 - - self.quant_method.apply( - layer=layer, - x=x, - router_logits=router_logits, - top_k=2, - renormalize=True, - global_num_experts=self.num_experts, - activation="gelu", - apply_router_weight_on_input=True, - mc2_mask=mc2_mask, - pertoken_scale=pertoken_scale, - ) - - fused_experts_input = mock_comm.fused_experts.call_args.kwargs["fused_experts_input"] - self.assertEqual(fused_experts_input.activation, "gelu") - self.assertTrue(fused_experts_input.routing.apply_router_weight_on_input) - self.assertIs(fused_experts_input.routing.mc2_mask, mc2_mask) - self.assertIs(fused_experts_input.routing.pertoken_scale, pertoken_scale) - self.assertIs(fused_experts_input.topk_weights, topk_weights) - self.assertIs(fused_experts_input.topk_ids, topk_ids) - - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_flash_common3_context") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts") - def test_apply_overlap_gate_uses_fc3_context( - self, - mock_select_experts, - mock_extra_ctx, - mock_get_flash_common3_context, - ): - tokens = 4 - hidden_size = self.hidden_size - layer = torch.nn.Module() - layer.w13_weight = torch.randint( - -8, - 8, - (self.num_experts, 2 * self.intermediate_size, hidden_size), - dtype=torch.int8, - ) - layer.w2_weight = torch.randint( - -8, - 8, - (self.num_experts, hidden_size, self.intermediate_size), - dtype=torch.int8, - ) - layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size, dtype=torch.float32) - layer.w2_weight_scale = torch.ones(self.num_experts, hidden_size, dtype=torch.float32) - - x = torch.randn(tokens, hidden_size, dtype=torch.float32) - router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32) - topk_weights = torch.randn(tokens, 2, dtype=torch.float32) - topk_ids = torch.randint(0, self.num_experts, (tokens, 2), dtype=torch.int64) - mc2_mask = torch.tensor([1, 0, 1, 0], dtype=torch.bool) - pertoken_scale = torch.randn(tokens, dtype=torch.float32) - - self.quant_method.multistream_overlap_gate = True - self.quant_method.in_dtype = torch.float32 - mock_get_flash_common3_context.return_value = Mock(topk_weights=topk_weights, topk_ids=topk_ids) - - mock_comm = Mock() - mock_comm.fused_experts.return_value = torch.randn(tokens, hidden_size, dtype=torch.float32) - mock_extra_ctx.moe_comm_method = mock_comm - mock_extra_ctx.moe_comm_type = MoECommType.ALLGATHER - - self.quant_method.apply( - layer=layer, - x=x, - router_logits=router_logits, - top_k=2, - renormalize=True, - global_num_experts=self.num_experts, - activation="gelu", - apply_router_weight_on_input=True, - mc2_mask=mc2_mask, - pertoken_scale=pertoken_scale, - ) - - mock_select_experts.assert_not_called() - fused_experts_input = mock_comm.fused_experts.call_args.kwargs["fused_experts_input"] - self.assertEqual(fused_experts_input.activation, "gelu") - self.assertTrue(fused_experts_input.routing.apply_router_weight_on_input) - self.assertIs(fused_experts_input.routing.mc2_mask, mc2_mask) - self.assertIs(fused_experts_input.routing.pertoken_scale, pertoken_scale) - self.assertIs(fused_experts_input.topk_weights, topk_weights) - self.assertIs(fused_experts_input.topk_ids, topk_ids) From 267c440fcf00d036b4cc4e9ce0042101403fddea Mon Sep 17 00:00:00 2001 From: Wang Kunpeng <1289706727@qq.com> Date: Fri, 24 Apr 2026 16:39:16 +0800 Subject: [PATCH 2/4] [Test]Add quantization test case Signed-off-by: Wang Kunpeng <1289706727@qq.com> --- .../ut/quantization/conftest_quantization.py | 22 +-- .../ut/quantization/methods/test_registry.py | 13 +- .../methods/test_w4a4_flatquant.py | 16 +- .../methods/test_w4a4_laos_dynamic.py | 11 +- .../quantization/methods/test_w4a4_mxfp4.py | 50 +++--- tests/ut/quantization/methods/test_w4a8.py | 18 +- tests/ut/quantization/methods/test_w8a16.py | 15 +- .../quantization/methods/test_w8a8_dynamic.py | 125 ++++++++------ .../quantization/methods/test_w8a8_mxfp8.py | 58 ++++--- .../quantization/methods/test_w8a8_pdmix.py | 22 ++- .../quantization/methods/test_w8a8_static.py | 43 +++-- .../test_compressed_tensors_config.py | 19 +-- tests/ut/quantization/test_method_adapters.py | 33 ++-- .../ut/quantization/test_modelslim_config.py | 157 ++++++++++-------- tests/ut/quantization/test_quant_parser.py | 22 +-- tests/ut/quantization/test_quant_type.py | 1 - tests/ut/quantization/test_utils.py | 1 - 17 files changed, 339 insertions(+), 287 deletions(-) diff --git a/tests/ut/quantization/conftest_quantization.py b/tests/ut/quantization/conftest_quantization.py index e5435de42038..7ae697c2b760 100644 --- a/tests/ut/quantization/conftest_quantization.py +++ b/tests/ut/quantization/conftest_quantization.py @@ -46,12 +46,10 @@ "observer_kwargs": {}, "strategy": "token", "symmetric": True, - "type": "int" + "type": "int", }, "output_activations": None, - "targets": [ - "Linear" - ], + "targets": ["Linear"], "weights": { "actorder": None, "block_structure": None, @@ -62,23 +60,23 @@ "observer_kwargs": {}, "strategy": "channel", "symmetric": True, - "type": "int" - } + "type": "int", + }, } }, "format": "int-quantized", "global_compression_ratio": None, - "ignore": [ - "lm_head" - ], + "ignore": ["lm_head"], "kv_cache_scheme": None, "quant_method": "compressed-tensors", "quantization_status": "compressed", } + def identity(*args): return args[0] + def create_mock_vllm_config( quant_description=None, model_dtype=torch.bfloat16, @@ -182,9 +180,11 @@ def create_mxfp_moe_layer( torch.randn(num_experts, hidden_size, intermediate_size).to(weight_dtype), requires_grad=False ) layer.w13_weight_scale = nn.Parameter( - torch.randint(0, 255, (num_experts, 2 * intermediate_size, hidden_size // group_size), dtype=scale_dtype), requires_grad=False + torch.randint(0, 255, (num_experts, 2 * intermediate_size, hidden_size // group_size), dtype=scale_dtype), + requires_grad=False, ) layer.w2_weight_scale = nn.Parameter( - torch.randint(0, 255, (num_experts, hidden_size, intermediate_size // group_size), dtype=scale_dtype), requires_grad=False + torch.randint(0, 255, (num_experts, hidden_size, intermediate_size // group_size), dtype=scale_dtype), + requires_grad=False, ) return layer diff --git a/tests/ut/quantization/methods/test_registry.py b/tests/ut/quantization/methods/test_registry.py index ecb99a05077c..5aa145bf2636 100644 --- a/tests/ut/quantization/methods/test_registry.py +++ b/tests/ut/quantization/methods/test_registry.py @@ -5,13 +5,12 @@ ) from vllm_ascend.quantization.methods.registry import ( _SCHEME_REGISTRY, - register_scheme, get_scheme_class, + register_scheme, ) class TestRegisterScheme(TestBase): - def test_register_scheme(self): @register_scheme("TEST_QUANT_TYPE", "linear") class TestLinearScheme(AscendLinearScheme): @@ -26,7 +25,6 @@ def apply(self, layer, x, bias=None, tp_rank=0): class TestGetSchemeClass(TestBase): - def test_get_scheme_class_existing_linear(self): cls = get_scheme_class("W8A8_DYNAMIC", "linear") self.assertIsNotNone(cls) @@ -51,6 +49,7 @@ def test_get_scheme_class_nonexistent_layer_type(self): def test_register_scheme_duplicate_raises(self): with self.assertRaises(ValueError): + @register_scheme("W8A8_DYNAMIC", "linear") class Duplicate: pass @@ -60,8 +59,7 @@ def test_all_linear_schemes_subclass_ascend_linear_scheme(self): if layer_type == "linear": self.assertTrue( issubclass(scheme_cls, AscendLinearScheme), - f"{scheme_cls.__name__} for {quant_type}/{layer_type} " - f"should be subclass of AscendLinearScheme", + f"{scheme_cls.__name__} for {quant_type}/{layer_type} should be subclass of AscendLinearScheme", ) def test_all_moe_schemes_subclass_ascend_moe_scheme(self): @@ -69,8 +67,7 @@ def test_all_moe_schemes_subclass_ascend_moe_scheme(self): if layer_type == "moe": self.assertTrue( issubclass(scheme_cls, AscendMoEScheme), - f"{scheme_cls.__name__} for {quant_type}/{layer_type} " - f"should be subclass of AscendMoEScheme", + f"{scheme_cls.__name__} for {quant_type}/{layer_type} should be subclass of AscendMoEScheme", ) def test_registry_not_empty(self): @@ -80,4 +77,4 @@ def test_registry_key_format(self): for key in _SCHEME_REGISTRY.keys(): self.assertEqual(len(key), 2) self.assertIsInstance(key[0], str) - self.assertIsInstance(key[1], str) \ No newline at end of file + self.assertIsInstance(key[1], str) diff --git a/tests/ut/quantization/methods/test_w4a4_flatquant.py b/tests/ut/quantization/methods/test_w4a4_flatquant.py index 04f6cc0d5662..f48b250d52d1 100644 --- a/tests/ut/quantization/methods/test_w4a4_flatquant.py +++ b/tests/ut/quantization/methods/test_w4a4_flatquant.py @@ -5,11 +5,11 @@ import torch.nn as nn from vllm_ascend.quantization.methods.w4a4_flatquant import ( + KRONECKER_QUANT_MAX_BATCH_SIZE, AscendW4A4FlatQuantDynamicLinearMethod, + batched_kronecker_quant, get_decompose_dim, pack_int4_weights, - KRONECKER_QUANT_MAX_BATCH_SIZE, - batched_kronecker_quant, ) @@ -53,7 +53,7 @@ def test_pack_int4_weights_npu_success(self, mock_torch_npu): mock_torch_npu.npu_convert_weight_to_int4pack.assert_called_once() self.assertTrue(torch.equal(result, mock_packed_tensor)) - @patch('vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu') + @patch("vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu") def test_large_batch_multiple_calls(self, mock_npu): batch_size = 50000 x = torch.randn(batch_size, 24, 32) @@ -61,8 +61,10 @@ def test_large_batch_multiple_calls(self, mock_npu): right_trans = torch.randn(32, 32) num_chunks = batch_size // KRONECKER_QUANT_MAX_BATCH_SIZE + 1 mock_returns = [ - (torch.randint(0, 255, (KRONECKER_QUANT_MAX_BATCH_SIZE, 24, 4), dtype=torch.int32), - torch.randn(KRONECKER_QUANT_MAX_BATCH_SIZE)) + ( + torch.randint(0, 255, (KRONECKER_QUANT_MAX_BATCH_SIZE, 24, 4), dtype=torch.int32), + torch.randn(KRONECKER_QUANT_MAX_BATCH_SIZE), + ) for _ in range(num_chunks - 1) ] last_chunk_size = batch_size - (num_chunks - 1) * KRONECKER_QUANT_MAX_BATCH_SIZE @@ -73,7 +75,7 @@ def test_large_batch_multiple_calls(self, mock_npu): result_x, result_scale = batched_kronecker_quant(x, left_trans, right_trans, 0.95) self.assertEqual(mock_npu.npu_kronecker_quant.call_count, num_chunks) - @patch('vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu') + @patch("vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu") def test_exact_max_batch_size(self, mock_npu): batch_size = KRONECKER_QUANT_MAX_BATCH_SIZE x = torch.randn(batch_size, 24, 32) @@ -81,7 +83,7 @@ def test_exact_max_batch_size(self, mock_npu): right_trans = torch.randn(32, 32) mock_npu.npu_kronecker_quant.return_value = ( torch.randint(0, 255, (batch_size, 24, 4), dtype=torch.int32), - torch.randn(batch_size, dtype=torch.float32) + torch.randn(batch_size, dtype=torch.float32), ) result_x, result_scale = batched_kronecker_quant(x, left_trans, right_trans, 0.95) mock_npu.npu_kronecker_quant.assert_called_once() diff --git a/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py b/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py index 0da7b5dbb0fc..584db5fdbdbb 100644 --- a/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py +++ b/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py @@ -1,12 +1,13 @@ from unittest.mock import MagicMock, patch + import torch import torch.nn as nn + from tests.ut.base import TestBase from vllm_ascend.quantization.methods.w4a4_laos_dynamic import AscendW4A4LaosDynamicLinearMethod class TestAscendW4A4LaosDynamicLinearMethod(TestBase): - def setUp(self): self.method = AscendW4A4LaosDynamicLinearMethod() @@ -43,7 +44,7 @@ def test_get_pergroup_param_various_layer_types(self): def test_apply_2d_input(self, mock_dyn_quant, mock_matmul): mock_dyn_quant.return_value = ( torch.randint(0, 15, (32, 128), dtype=torch.int32), - torch.randn(32, dtype=torch.float32) + torch.randn(32, dtype=torch.float32), ) mock_matmul.return_value = torch.randn(32, 256) layer = MagicMock() @@ -58,7 +59,7 @@ def test_apply_2d_input(self, mock_dyn_quant, mock_matmul): def test_apply_with_bias(self, mock_dyn_quant, mock_matmul): mock_dyn_quant.return_value = ( torch.randint(0, 15, (32, 128), dtype=torch.int32), - torch.randn(32, dtype=torch.float32) + torch.randn(32, dtype=torch.float32), ) expected_output = torch.randn(32, 256, dtype=torch.bfloat16) mock_matmul.return_value = expected_output @@ -86,7 +87,9 @@ def test_process_weights_various_input_sizes(self, mock_convert): for input_size, output_size in [(64, 128), (256, 512)]: mock_convert.return_value = torch.randint(0, 15, (output_size, input_size // 8), dtype=torch.int32) layer = nn.Module() - layer.weight = nn.Parameter(torch.randint(-8, 7, (output_size, input_size), dtype=torch.int8), requires_grad=False) + layer.weight = nn.Parameter( + torch.randint(-8, 7, (output_size, input_size), dtype=torch.int8), requires_grad=False + ) layer.weight_scale = nn.Parameter(torch.randn(output_size, 1, dtype=torch.float32), requires_grad=False) self.method.process_weights_after_loading(layer) mock_convert.assert_called() diff --git a/tests/ut/quantization/methods/test_w4a4_mxfp4.py b/tests/ut/quantization/methods/test_w4a4_mxfp4.py index 3f3da32ffe26..d594cd04c5cb 100644 --- a/tests/ut/quantization/methods/test_w4a4_mxfp4.py +++ b/tests/ut/quantization/methods/test_w4a4_mxfp4.py @@ -1,18 +1,18 @@ 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_vllm_config, create_mock_ascend_config -from vllm_ascend.device.mxfp_compat import FLOAT8_E8M0FNU_DTYPE +from tests.ut.quantization.conftest_quantization import create_mock_ascend_config, create_mock_vllm_config +from vllm_ascend.quantization.methods.base import QuantType from vllm_ascend.quantization.methods.w4a4_mxfp4 import ( - AscendW4A4MXFP4DynamicLinearMethod, AscendW4A4MXFP4DynamicFusedMoEMethod, + AscendW4A4MXFP4DynamicLinearMethod, ) -from vllm_ascend.quantization.methods.base import QuantType 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): @@ -64,14 +64,14 @@ def test_group_size_from_config(self, mock_vllm_config, mock_ensure): 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) + 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): + with patch.object(self.scheme, "group_size", 32): output = self.scheme.apply(layer, x) self.assertEqual(output.shape[0], 32) @@ -100,16 +100,18 @@ 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["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) + 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) @@ -118,8 +120,9 @@ 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.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)) @@ -135,10 +138,12 @@ def test_apply_full_params(self, mock_select, mock_ctx, mock_npu): 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) + 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) @@ -149,9 +154,14 @@ def test_apply_full_params(self, mock_select, mock_ctx, mock_npu): mock_ctx.moe_comm_method = mock_comm mock_ctx.moe_comm_type = Mock() result = self.scheme.apply( - layer, x, router_logits, top_k=2, renormalize=True, + layer, + x, + router_logits, + top_k=2, + renormalize=True, global_num_experts=self.num_experts, - activation="silu", pertoken_scale=torch.randn(tokens), - apply_router_weight_on_input=True + activation="silu", + pertoken_scale=torch.randn(tokens), + apply_router_weight_on_input=True, ) mock_comm.fused_experts.assert_called_once() diff --git a/tests/ut/quantization/methods/test_w4a8.py b/tests/ut/quantization/methods/test_w4a8.py index 0b75952d8ea6..374f4b2b954c 100644 --- a/tests/ut/quantization/methods/test_w4a8.py +++ b/tests/ut/quantization/methods/test_w4a8.py @@ -1,4 +1,4 @@ -from unittest.mock import Mock, patch, MagicMock +from unittest.mock import MagicMock, Mock, patch import torch @@ -269,14 +269,16 @@ def test_get_weight_compressed_tensors(self): 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) + 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('torch_npu.npu_quantize') - @patch('torch.Tensor.npu') + @patch("torch_npu.npu_quantize") + @patch("torch.Tensor.npu") def test_process_weights_after_loading_compressed_tensors(self, mock_npu, mock_npu_quantize): mock_npu.return_value = torch.Tensor() mock_npu_quantize.return_value = torch.Tensor() @@ -286,12 +288,10 @@ def test_process_weights_after_loading_compressed_tensors(self, mock_npu, mock_n 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.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.shape, (self.experts, self.output_size)) self.assertEqual(layer.w2_scale_bias.data.dtype, torch.float32) @patch("vllm_ascend.quantization.methods.w4a8._EXTRA_CTX") @@ -368,4 +368,4 @@ def test_apply_comprehensive(self, mock_build_input, mock_select, mock_ctx): 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)) \ No newline at end of file + self.assertTrue(torch.equal(output, expected_output)) diff --git a/tests/ut/quantization/methods/test_w8a16.py b/tests/ut/quantization/methods/test_w8a16.py index 916342a3fe34..755de13e4dc3 100644 --- a/tests/ut/quantization/methods/test_w8a16.py +++ b/tests/ut/quantization/methods/test_w8a16.py @@ -16,20 +16,20 @@ def test_get_weight(self): sizes = [(64, 128), (256, 512), (1024, 2048), (1, 1)] for input_size, output_size in sizes: weight = self.method.get_weight(input_size, output_size) - self.assertEqual(weight['weight'].dtype, torch.int8) - self.assertEqual(weight['weight'].shape, (output_size, input_size)) + self.assertEqual(weight["weight"].dtype, torch.int8) + self.assertEqual(weight["weight"].shape, (output_size, input_size)) self.assertEqual(len(weight), 1) weight = self.method.get_weight(256, 128, torch.float16) - self.assertEqual(weight['weight'].dtype, torch.int8) + self.assertEqual(weight["weight"].dtype, torch.int8) def test_get_per_channel_param(self): for output_size, dtype in [(128, torch.bfloat16), (256, torch.float16)]: per_channel_params = self.method.get_perchannel_param(output_size, dtype) - self.assertEqual(per_channel_params['weight_scale'].dtype, dtype) - self.assertEqual(per_channel_params['weight_scale'].shape, (output_size, 1)) - self.assertEqual(per_channel_params['weight_offset'].dtype, dtype) - self.assertEqual(per_channel_params['weight_offset'].shape, (output_size, 1)) + self.assertEqual(per_channel_params["weight_scale"].dtype, dtype) + self.assertEqual(per_channel_params["weight_scale"].shape, (output_size, 1)) + self.assertEqual(per_channel_params["weight_offset"].dtype, dtype) + self.assertEqual(per_channel_params["weight_offset"].shape, (output_size, 1)) self.assertEqual(len(per_channel_params), 2) def test_get_pertensor_param_returns_empty(self): @@ -40,7 +40,6 @@ def test_get_pergroup_param_returns_empty(self): result = self.method.get_pergroup_param(256, 128, torch.bfloat16) self.assertEqual(result, {}) - @patch("torch_npu.npu_weight_quant_batchmatmul") def test_apply_with_x_is_int8(self, mock_npu_weight_quant_batchmatmul): layer = MagicMock() diff --git a/tests/ut/quantization/methods/test_w8a8_dynamic.py b/tests/ut/quantization/methods/test_w8a8_dynamic.py index 0fa6ed9d08e4..e84e2bb1a503 100644 --- a/tests/ut/quantization/methods/test_w8a8_dynamic.py +++ b/tests/ut/quantization/methods/test_w8a8_dynamic.py @@ -1,38 +1,37 @@ -from unittest.mock import Mock, patch, MagicMock +from unittest.mock import MagicMock, Mock, patch import torch from tests.ut.base import TestBase -from tests.ut.quantization.conftest_quantization import create_mock_vllm_config, create_mock_ascend_config, \ - create_moe_layer +from tests.ut.quantization.conftest_quantization import ( + create_mock_ascend_config, + create_mock_vllm_config, + create_moe_layer, +) from vllm_ascend.ascend_forward_context import MoECommType from vllm_ascend.quantization.methods.w8a8_dynamic import ( - AscendW8A8DynamicLinearMethod, AscendW8A8DynamicFusedMoEMethod, + AscendW8A8DynamicLinearMethod, scale_from_float_to_int64, ) from vllm_ascend.quantization.quant_type import QuantType class TestScaleFromFloatToInt64(TestBase): - def test_scale_conversion_basic(self): scale = torch.tensor([0.5], dtype=torch.float32) - with patch.object(scale, 'cpu', return_value=scale): - with patch.object(scale, 'to', return_value=scale): - result = scale_from_float_to_int64(scale) + with patch.object(scale, "cpu", return_value=scale), patch.object(scale, "to", return_value=scale): + result = scale_from_float_to_int64(scale) self.assertEqual(result.dtype, torch.int64) def test_scale_conversion_preserves_device(self): scale = torch.tensor([0.5, 1.0], dtype=torch.float32) - with patch.object(scale, 'cpu', return_value=scale): - with patch.object(scale, 'to', return_value=scale): - result = scale_from_float_to_int64(scale) + with patch.object(scale, "cpu", return_value=scale), patch.object(scale, "to", return_value=scale): + result = scale_from_float_to_int64(scale) self.assertEqual(result.shape, (2,)) class TestAscendW8A8DynamicLinearMethod(TestBase): - def setUp(self): self.method = AscendW8A8DynamicLinearMethod() @@ -40,24 +39,24 @@ def test_get_weight_various_sizes(self): sizes = [(64, 128), (256, 512), (1024, 2048)] for input_size, output_size in sizes: weight = self.method.get_weight(input_size, output_size, torch.bfloat16) - self.assertEqual(weight['weight'].dtype, torch.int8) - self.assertEqual(weight['weight'].shape, (output_size, input_size)) + self.assertEqual(weight["weight"].dtype, torch.int8) + self.assertEqual(weight["weight"].shape, (output_size, input_size)) def test_get_perchannel_param_dtype_variations(self): dtypes = [torch.bfloat16, torch.float16] for dtype in dtypes: params = self.method.get_perchannel_param(128, dtype) - self.assertEqual(params['weight_scale'].dtype, dtype) - self.assertEqual(params['weight_offset'].dtype, dtype) - self.assertEqual(params['weight_scale'].shape, (128, 1)) - self.assertEqual(params['weight_offset'].shape, (128, 1)) + self.assertEqual(params["weight_scale"].dtype, dtype) + self.assertEqual(params["weight_offset"].dtype, dtype) + self.assertEqual(params["weight_scale"].shape, (128, 1)) + self.assertEqual(params["weight_offset"].shape, (128, 1)) @patch("torch_npu.npu_quant_matmul") @patch("torch_npu.npu_dynamic_quant") def test_apply_2d_input(self, mock_dyn_quant, mock_matmul): mock_dyn_quant.return_value = ( torch.randint(-128, 127, (32, 128), dtype=torch.int8), - torch.randn(32, dtype=torch.float32) + torch.randn(32, dtype=torch.float32), ) mock_matmul.return_value = torch.randn(32, 256) layer = MagicMock() @@ -73,7 +72,7 @@ def test_apply_2d_input(self, mock_dyn_quant, mock_matmul): def test_apply_3d_input_with_squeeze(self, mock_dyn_quant, mock_matmul): mock_dyn_quant.return_value = ( torch.randint(-128, 127, (32, 1, 128), dtype=torch.int8), - torch.randn(32, 1, dtype=torch.float32) + torch.randn(32, 1, dtype=torch.float32), ) mock_matmul.return_value = torch.randn(32, 1, 256) layer = MagicMock() @@ -86,9 +85,9 @@ def test_apply_3d_input_with_squeeze(self, mock_dyn_quant, mock_matmul): def test_process_weights_after_loading(self): layer = MagicMock() layer.weight.data = torch.randint(-128, 127, (128, 256), dtype=torch.int8) - layer.weight_scale.data=torch.randn(256, 1, dtype=torch.bfloat16) - layer.weight_offset.data=torch.randn(256, 1, dtype=torch.bfloat16) - with patch('vllm_ascend.quantization.methods.w8a8_dynamic.maybe_trans_nz', side_effect=lambda x: x): + layer.weight_scale.data = torch.randn(256, 1, dtype=torch.bfloat16) + layer.weight_offset.data = torch.randn(256, 1, dtype=torch.bfloat16) + with patch("vllm_ascend.quantization.methods.w8a8_dynamic.maybe_trans_nz", side_effect=lambda x: x): self.method.process_weights_after_loading(layer) self.assertEqual(layer.weight_scale_fp32.dtype, torch.float32) self.assertEqual(layer.weight_scale.data.shape, (256,)) @@ -106,12 +105,15 @@ class TestAscendW8A8FusedMoEMethod(TestBase): @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ascend_config") @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ep_group") def setUp(self, mock_ep, mock_ascend, mock_mc2, mock_rank): - with patch('vllm_ascend.quantization.methods.w8a8_dynamic.get_current_vllm_config') as mock_vllm: + with patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_current_vllm_config") as mock_vllm: mock_vllm.return_value = create_mock_vllm_config() mock_ep.return_value = Mock() mock_ascend.return_value = create_mock_ascend_config() - mock_mc2.return_value = MagicMock(device_group=Mock( - _get_backend=Mock(return_value=Mock(get_hccl_comm_name=Mock(return_value="test_comm"))))) + mock_mc2.return_value = MagicMock( + device_group=Mock( + _get_backend=Mock(return_value=Mock(get_hccl_comm_name=Mock(return_value="test_comm"))) + ) + ) mock_rank.return_value = 0 self.quant_method = AscendW8A8DynamicFusedMoEMethod() @@ -121,14 +123,16 @@ def test_quant_type_is_w8a8(self): def test_get_weight_various_expert_counts(self): expert_counts = [4, 8, 16, 32] for num_experts in expert_counts: - param_dict = self.quant_method.get_weight(num_experts, self.intermediate_size, self.hidden_size, - torch.bfloat16) + param_dict = self.quant_method.get_weight( + num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16 + ) self.assertEqual(param_dict["w13_weight"].shape[0], num_experts) self.assertEqual(param_dict["w2_weight"].shape[0], num_experts) def test_get_dynamic_quant_param_various_sizes(self): - param_dict = self.quant_method.get_dynamic_quant_param(self.num_experts, self.intermediate_size, - self.hidden_size, torch.bfloat16) + param_dict = self.quant_method.get_dynamic_quant_param( + self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16 + ) self.assertEqual(param_dict["w13_weight_scale"].dtype, torch.bfloat16) self.assertEqual(param_dict["w13_weight_offset"].shape, (self.num_experts, 2 * self.intermediate_size, 1)) self.assertEqual(param_dict["w2_weight_scale"].dtype, torch.bfloat16) @@ -262,10 +266,12 @@ def test_apply_overlap_gate_uses_fc3_context( def test_apply_with_zero_experts(self, mock_zero, mock_select, mock_ctx): tokens = 4 layer = MagicMock() - layer.w13_weight = torch.randint(-8, 8, (self.num_experts, 2 * self.intermediate_size, self.hidden_size), - dtype=torch.int8) - layer.w2_weight = torch.randint(-8, 8, (self.num_experts, self.hidden_size, self.intermediate_size), - dtype=torch.int8) + layer.w13_weight = torch.randint( + -8, 8, (self.num_experts, 2 * self.intermediate_size, self.hidden_size), dtype=torch.int8 + ) + layer.w2_weight = torch.randint( + -8, 8, (self.num_experts, self.hidden_size, self.intermediate_size), dtype=torch.int8 + ) layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size) layer.w2_weight_scale = torch.ones(self.num_experts, self.hidden_size) layer.zero_expert_num = 2 @@ -290,10 +296,12 @@ def test_apply_with_zero_experts(self, mock_zero, mock_select, mock_ctx): def test_apply_with_enable_force_load_balance(self, mock_select, mock_ctx): tokens = 4 layer = MagicMock() - layer.w13_weight = torch.randint(-8, 8, (self.num_experts, 2 * self.intermediate_size, self.hidden_size), - dtype=torch.int8) - layer.w2_weight = torch.randint(-8, 8, (self.num_experts, self.hidden_size, self.intermediate_size), - dtype=torch.int8) + layer.w13_weight = torch.randint( + -8, 8, (self.num_experts, 2 * self.intermediate_size, self.hidden_size), dtype=torch.int8 + ) + layer.w2_weight = torch.randint( + -8, 8, (self.num_experts, self.hidden_size, self.intermediate_size), dtype=torch.int8 + ) layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size) layer.w2_weight_scale = torch.ones(self.num_experts, self.hidden_size) layer.zero_expert_num = 0 @@ -310,30 +318,39 @@ def test_apply_with_enable_force_load_balance(self, mock_select, mock_ctx): mock_ctx.moe_comm_method = mock_comm mock_ctx.moe_comm_type = Mock() self.quant_method.in_dtype = torch.float32 - self.quant_method.apply(layer, x, router_logits, top_k=2, renormalize=True, global_num_experts=self.num_experts, - enable_force_load_balance=True) + self.quant_method.apply( + layer, + x, + router_logits, + top_k=2, + renormalize=True, + global_num_experts=self.num_experts, + enable_force_load_balance=True, + ) - @patch('torch_npu.npu_format_cast') + @patch("torch_npu.npu_format_cast") @patch("vllm_ascend.quantization.methods.w8a8_dynamic.envs_ascend") def test_process_weights_with_fused_mc2(self, mock_envs, mock_format_cast): mock_envs.VLLM_ASCEND_ENABLE_FUSED_MC2 = 1 - mock_format_cast.return_value = torch.randint(-8, 8, - (self.num_experts, self.hidden_size, 2 * self.intermediate_size), - dtype=torch.int8) - layer = create_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size) + mock_format_cast.return_value = torch.randint( + -8, 8, (self.num_experts, self.hidden_size, 2 * self.intermediate_size), dtype=torch.int8 + ) + layer = create_moe_layer( + num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size + ) self.quant_method.process_weights_after_loading(layer) - self.assertTrue(hasattr(layer, 'w13_weight_scale_fp32')) + self.assertTrue(hasattr(layer, "w13_weight_scale_fp32")) - @patch('torch_npu.npu_format_cast') + @patch("torch_npu.npu_format_cast") @patch("vllm_ascend.quantization.methods.w8a8_dynamic.envs_ascend") def test_process_weights_with_dynamic_eplb(self, mock_envs, mock_format_cast): mock_envs.VLLM_ASCEND_ENABLE_FUSED_MC2 = 0 self.quant_method.dynamic_eplb = True - mock_format_cast.return_value = torch.randint(-8, 8, - (self.num_experts, self.hidden_size, 2 * self.intermediate_size), - dtype=torch.int8) - layer = create_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size) + mock_format_cast.return_value = torch.randint( + -8, 8, (self.num_experts, self.hidden_size, 2 * self.intermediate_size), dtype=torch.int8 + ) + layer = create_moe_layer( + num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size + ) self.quant_method.process_weights_after_loading(layer) - self.assertTrue(hasattr(layer, 'w13_weight_list')) + self.assertTrue(hasattr(layer, "w13_weight_list")) diff --git a/tests/ut/quantization/methods/test_w8a8_mxfp8.py b/tests/ut/quantization/methods/test_w8a8_mxfp8.py index f3a740bcfd25..81188414b612 100644 --- a/tests/ut/quantization/methods/test_w8a8_mxfp8.py +++ b/tests/ut/quantization/methods/test_w8a8_mxfp8.py @@ -5,19 +5,18 @@ from tests.ut.base import TestBase from tests.ut.quantization.conftest_quantization import ( - create_mock_vllm_config, create_mock_ascend_config, + create_mock_vllm_config, create_mxfp_moe_layer, ) from vllm_ascend.quantization.methods.base import QuantType from vllm_ascend.quantization.methods.w8a8_mxfp8 import ( - AscendW8A8MXFP8DynamicLinearMethod, AscendW8A8MXFP8DynamicFusedMoEMethod, + AscendW8A8MXFP8DynamicLinearMethod, ) class TestAscendW8A8MXFP8LinearMethod(TestBase): - @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.ensure_mxfp8_linear_available") @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.get_current_vllm_config") def setUp(self, mock_vllm, mock_ensure): @@ -104,10 +103,11 @@ def test_scale_shape_transformation(self): @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.torch_npu") def test_apply(self, mock_torch_npu): from vllm_ascend.device.mxfp_compat import FLOAT8_E8M0FNU_DTYPE + dynamic_scale = torch.randint(0, 255, (32, 8), dtype=torch.uint8) mock_torch_npu.npu_dynamic_mx_quant.return_value = ( torch.randint(0, 255, (32, 256), dtype=torch.uint8), - dynamic_scale + dynamic_scale, ) mock_torch_npu.npu_quant_matmul.return_value = torch.randn(32, 128, dtype=torch.float16) layer = nn.Module() @@ -118,15 +118,14 @@ def test_apply(self, mock_torch_npu): output = self.scheme.apply(layer, x, bias) self.assertEqual(output.shape, (32, 1, 128)) call_kwargs = mock_torch_npu.npu_quant_matmul.call_args.kwargs - self.assertEqual(call_kwargs['bias'].dtype, torch.float32) - self.assertEqual(call_kwargs['group_sizes'], [1, 1, self.scheme.group_size]) - self.assertEqual(call_kwargs['scale_dtype'], FLOAT8_E8M0FNU_DTYPE) - self.assertEqual(call_kwargs['pertoken_scale_dtype'], FLOAT8_E8M0FNU_DTYPE) - self.assertEqual(call_kwargs['output_dtype'], torch.float16) + self.assertEqual(call_kwargs["bias"].dtype, torch.float32) + self.assertEqual(call_kwargs["group_sizes"], [1, 1, self.scheme.group_size]) + self.assertEqual(call_kwargs["scale_dtype"], FLOAT8_E8M0FNU_DTYPE) + self.assertEqual(call_kwargs["pertoken_scale_dtype"], FLOAT8_E8M0FNU_DTYPE) + self.assertEqual(call_kwargs["output_dtype"], torch.float16) class TestAscendW8A8MXFP8MoEMethod(TestBase): - num_experts = 8 hidden_size = 128 intermediate_size = 256 @@ -159,25 +158,33 @@ def test_get_weight_various_expert_counts(self): self.assertEqual(result["w13_weight"].shape[0], num_experts) def test_get_dynamic_quant_param_dtype_uint8(self): - result = self.scheme.get_dynamic_quant_param(self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16) + result = self.scheme.get_dynamic_quant_param( + self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16 + ) self.assertEqual(result["w13_weight_scale"].dtype, torch.uint8) self.assertEqual(result["w2_weight_scale"].dtype, torch.uint8) def test_process_weights_stores_original_shapes(self): - layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + layer = create_mxfp_moe_layer( + num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size + ) self.scheme.process_weights_after_loading(layer) self.assertTrue(hasattr(layer, "_mxfp8_original_shapes")) self.assertIn("w13_weight", layer._mxfp8_original_shapes) def test_process_weights_double_call_idempotent(self): - layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + layer = create_mxfp_moe_layer( + num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size + ) self.scheme.process_weights_after_loading(layer) weight_after = layer.w13_weight.data.clone() self.scheme.process_weights_after_loading(layer) self.assertTrue(torch.equal(layer.w13_weight.data, weight_after)) def test_restore_weights_for_rl_loading(self): - layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + layer = create_mxfp_moe_layer( + num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size + ) original_w13_shape = layer.w13_weight.shape self.scheme.process_weights_after_loading(layer) self.assertNotEqual(layer.w13_weight.shape, original_w13_shape) @@ -185,13 +192,17 @@ def test_restore_weights_for_rl_loading(self): self.assertEqual(layer.w13_weight.shape, original_w13_shape) def test_restore_without_transform_is_noop(self): - layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + layer = create_mxfp_moe_layer( + num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size + ) original_weight = layer.w13_weight.data.clone() self.scheme.restore_weights_for_rl_loading(layer) self.assertTrue(torch.equal(layer.w13_weight.data, original_weight)) def test_process_weights_transposes_weights(self): - layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + layer = create_mxfp_moe_layer( + num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size + ) original_shape = layer.w13_weight.shape self.scheme.process_weights_after_loading(layer) self.assertEqual(layer.w13_weight.shape, (original_shape[0], original_shape[2], original_shape[1])) @@ -200,7 +211,9 @@ def test_process_weights_transposes_weights(self): @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.select_experts") def test_apply_full_params(self, mock_select, mock_ctx): tokens = 4 - layer = create_mxfp_moe_layer(num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size) + layer = create_mxfp_moe_layer( + num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size + ) self.scheme.process_weights_after_loading(layer) x = torch.randn(tokens, self.hidden_size, dtype=torch.bfloat16) router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32) @@ -212,9 +225,14 @@ def test_apply_full_params(self, mock_select, mock_ctx): mock_ctx.moe_comm_method = mock_comm mock_ctx.moe_comm_type = Mock() result = self.scheme.apply( - layer, x, router_logits, top_k=2, renormalize=True, + layer, + x, + router_logits, + top_k=2, + renormalize=True, global_num_experts=self.num_experts, - activation="silu", pertoken_scale=torch.randn(tokens) + activation="silu", + pertoken_scale=torch.randn(tokens), ) mock_select.assert_called_once() - mock_comm.fused_experts.assert_called_once() \ No newline at end of file + mock_comm.fused_experts.assert_called_once() diff --git a/tests/ut/quantization/methods/test_w8a8_pdmix.py b/tests/ut/quantization/methods/test_w8a8_pdmix.py index 723b8ec39b3a..fe4fe9f763be 100644 --- a/tests/ut/quantization/methods/test_w8a8_pdmix.py +++ b/tests/ut/quantization/methods/test_w8a8_pdmix.py @@ -1,9 +1,15 @@ from unittest.mock import MagicMock, patch + import torch + from tests.ut.base import TestBase from tests.ut.quantization.conftest_quantization import create_mock_vllm_config -from vllm_ascend.quantization.methods import AscendW8A8LinearMethod, AscendW8A8DynamicLinearMethod, \ - AscendW8A8PDMixLinearMethod, AscendW8A8PDMixFusedMoeMethod +from vllm_ascend.quantization.methods import ( + AscendW8A8DynamicLinearMethod, + AscendW8A8LinearMethod, + AscendW8A8PDMixFusedMoeMethod, + AscendW8A8PDMixLinearMethod, +) class TestAscendW8A8PDMixLinearScheme(TestBase): @@ -34,7 +40,6 @@ def test_get_weight_delegates_to_static(self, mock_vllm_config, mock_dynamic_cls mock_static_instance.get_weight.assert_called_with(input_size, output_size, torch.bfloat16) mock_dynamic_instance.get_weight.assert_not_called() - @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") @@ -104,7 +109,9 @@ def test_apply_uses_dynamic_for_non_kv_consumer(self, mock_vllm_config, mock_dyn @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") - def test_process_weights_after_loading_sets_is_kv_consumer(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls): + def test_process_weights_after_loading_sets_is_kv_consumer( + self, mock_vllm_config, mock_dynamic_cls, mock_static_cls + ): mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None) mock_static_instance = MagicMock() mock_static_cls.return_value = mock_static_instance @@ -135,7 +142,6 @@ def test_is_kv_consumer_when_kv_transfer_config_set(self, mock_vllm_config, mock class TestAscendW8A8PDMixMoEScheme(TestBase): - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_mc2_group") @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ep_group") @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_current_vllm_config") @@ -146,7 +152,9 @@ def test_get_dynamic_quant_param(self, mock_ascend, mock_vllm, mock_ep, mock_mc2 mock_ascend.return_value = MagicMock(eplb_config=MagicMock(dynamic_eplb=False)) scheme = AscendW8A8PDMixFusedMoeMethod() num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype = 8, 256, 128, torch.bfloat16 - result = scheme.get_dynamic_quant_param(num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype) + result = scheme.get_dynamic_quant_param( + num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype + ) # test adds extra params self.assertEqual(result["w2_deq_scale"].shape, (num_experts, hidden_sizes)) self.assertEqual(result["w2_deq_scale"].dtype, torch.float32) @@ -161,4 +169,4 @@ def test_get_dynamic_quant_param(self, mock_ascend, mock_vllm, mock_ep, mock_mc2 self.assertIn("w13_weight_scale", result) self.assertIn("w13_weight_offset", result) self.assertIn("w2_weight_scale", result) - self.assertIn("w2_weight_offset", result) \ No newline at end of file + self.assertIn("w2_weight_offset", result) diff --git a/tests/ut/quantization/methods/test_w8a8_static.py b/tests/ut/quantization/methods/test_w8a8_static.py index 98f8e0c9c7a1..5c3e9ecaab82 100644 --- a/tests/ut/quantization/methods/test_w8a8_static.py +++ b/tests/ut/quantization/methods/test_w8a8_static.py @@ -6,7 +6,7 @@ from tests.ut.base import TestBase from tests.ut.quantization.conftest_quantization import identity from vllm_ascend.quantization.methods.w8a8_static import AscendW8A8LinearMethod -from vllm_ascend.utils import AscendDeviceType, COMPRESSED_TENSORS_METHOD +from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD class TestAscendW8A8LinearMethod(TestBase): @@ -17,36 +17,36 @@ def test_get_weight(self): sizes = [(64, 128), (256, 512), (1024, 2048), (1, 1)] for input_size, output_size in sizes: weight = self.method.get_weight(input_size, output_size) - self.assertEqual(weight['weight'].dtype, torch.int8) - self.assertEqual(weight['weight'].shape, (output_size, input_size)) + self.assertEqual(weight["weight"].dtype, torch.int8) + self.assertEqual(weight["weight"].shape, (output_size, input_size)) self.assertEqual(len(weight), 1) weight = self.method.get_weight(256, 128, torch.float16) - self.assertEqual(weight['weight'].dtype, torch.int8) + self.assertEqual(weight["weight"].dtype, torch.int8) def test_get_pertensor_param(self): dtypes = [torch.bfloat16, torch.float16, torch.float32] for dtype in dtypes: params = self.method.get_pertensor_param(dtype) - self.assertEqual(params['input_scale'].dtype, dtype) - self.assertEqual(params['input_offset'].dtype, torch.int8) - self.assertEqual(params['input_scale'].shape, (1,)) - self.assertEqual(params['input_offset'].shape, (1,)) + self.assertEqual(params["input_scale"].dtype, dtype) + self.assertEqual(params["input_offset"].dtype, torch.int8) + self.assertEqual(params["input_scale"].shape, (1,)) + self.assertEqual(params["input_offset"].shape, (1,)) def test_get_perchannel_param(self): for output_size, dtype in [(128, torch.bfloat16), (256, torch.float16)]: params = self.method.get_perchannel_param(output_size, dtype) - self.assertEqual(params['quant_bias'].shape, (output_size,)) - self.assertEqual(params['quant_bias'].dtype, torch.int32) - self.assertEqual(params['weight_scale'].shape, (output_size, 1)) - self.assertEqual(params['weight_scale'].dtype, dtype) - self.assertEqual(params['weight_offset'].shape, (output_size, 1)) - self.assertEqual(params['weight_offset'].dtype, dtype) - self.assertEqual(params['deq_scale'].shape, (output_size, )) + self.assertEqual(params["quant_bias"].shape, (output_size,)) + self.assertEqual(params["quant_bias"].dtype, torch.int32) + self.assertEqual(params["weight_scale"].shape, (output_size, 1)) + self.assertEqual(params["weight_scale"].dtype, dtype) + self.assertEqual(params["weight_offset"].shape, (output_size, 1)) + self.assertEqual(params["weight_offset"].dtype, dtype) + self.assertEqual(params["deq_scale"].shape, (output_size,)) if dtype == torch.bfloat16: - self.assertEqual(params['deq_scale'].dtype, torch.float32) + self.assertEqual(params["deq_scale"].dtype, torch.float32) elif dtype == torch.float16: - self.assertEqual(params['deq_scale'].dtype, torch.int64) + self.assertEqual(params["deq_scale"].dtype, torch.int64) @patch("vllm_ascend.quantization.methods.w8a8_static.get_weight_prefetch_method") @patch("torch.ops.vllm.quantize") @@ -75,7 +75,7 @@ def test_apply_with_x_not_int8(self, mock_npu_quant_matmul, mock_quantize, mock_ mock_quantize.assert_called_once() mock_npu_quant_matmul.assert_called_once() call_kwargs = mock_npu_quant_matmul.call_args.kwargs - self.assertTrue(torch.equal(call_kwargs['bias'], quant_bias)) + self.assertTrue(torch.equal(call_kwargs["bias"], quant_bias)) @patch("torch.ops.vllm.quantize") @patch("torch_npu.npu_quant_matmul") @@ -98,7 +98,7 @@ def test_apply_with_x_is_int8(self, mock_npu_quant_matmul, mock_quantize): mock_quantize.assert_not_called() mock_npu_quant_matmul.assert_called_once() call_kwargs = mock_npu_quant_matmul.call_args.kwargs - self.assertTrue(torch.equal(call_kwargs['bias'], bias)) + self.assertTrue(torch.equal(call_kwargs["bias"], bias)) @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "0"}) @patch("torch_npu.npu_format_cast") @@ -149,9 +149,8 @@ def test_process_weights_after_loading_with_nz1(self, mock_npu_format_cast): self.assertTrue(isinstance(layer.deq_scale, MagicMock)) @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "2"}) - @patch('torch_npu.npu_format_cast') - def test_process_weights_after_loading_with_nz2_and_compressed_tensors(self, - mock_npu_format_cast): + @patch("torch_npu.npu_format_cast") + def test_process_weights_after_loading_with_nz2_and_compressed_tensors(self, mock_npu_format_cast): layer = MagicMock() layer.weight.data = torch.randint(-128, 127, (128, 256), dtype=torch.int8) diff --git a/tests/ut/quantization/test_compressed_tensors_config.py b/tests/ut/quantization/test_compressed_tensors_config.py index 28db27fd1bb6..b24a660924db 100644 --- a/tests/ut/quantization/test_compressed_tensors_config.py +++ b/tests/ut/quantization/test_compressed_tensors_config.py @@ -1,21 +1,20 @@ from unittest.mock import MagicMock, patch import torch - -from tests.ut.base import TestBase -from tests.ut.quantization.conftest_quantization import COMPRESSED_TENSORS_W8A8_CONFIG from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.linear import RowParallelLinear, UnquantizedLinearMethod + +from tests.ut.base import TestBase +from tests.ut.quantization.conftest_quantization import COMPRESSED_TENSORS_W8A8_CONFIG from vllm_ascend.ops.fused_moe.fused_moe import AscendUnquantizedFusedMoEMethod from vllm_ascend.quantization.compressed_tensors_config import AscendCompressedTensorsConfig -from vllm_ascend.quantization.method_adapters import AscendLinearMethod, AscendFusedMoEMethod -from vllm_ascend.quantization.methods import AscendW8A8DynamicLinearMethod, AscendW8A8DynamicFusedMoEMethod +from vllm_ascend.quantization.method_adapters import AscendFusedMoEMethod, AscendLinearMethod +from vllm_ascend.quantization.methods import AscendW8A8DynamicFusedMoEMethod, AscendW8A8DynamicLinearMethod from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD class TestAscendCompressedTensorsConfigBasic(TestBase): - def test_get_name(self): config = AscendCompressedTensorsConfig.from_config(COMPRESSED_TENSORS_W8A8_CONFIG) self.assertEqual(config.get_name(), "compressed-tensors") @@ -56,7 +55,6 @@ def test_apply_vllm_mapper(self): class TestAscendCompressedTensorsQuanType(TestBase): - def setUp(self): self.config = AscendCompressedTensorsConfig( target_scheme_map={"Linear": {}}, @@ -102,6 +100,7 @@ def test_detect_w4a8_dynamic(self): def test_detect_w4a16(self): from compressed_tensors.quantization import QuantizationType + weight = MagicMock() weight.num_bits = 4 weight.strategy = "group" @@ -118,7 +117,6 @@ def test_detect_unsupported_raises(self): class TestAscendCompressedTensorsConfigGetQuantMethod(TestBase): - def setUp(self): self.config = AscendCompressedTensorsConfig.from_config(COMPRESSED_TENSORS_W8A8_CONFIG) @@ -135,7 +133,8 @@ def test_get_linear_unquantized_method(self): self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) self.assertTrue(isinstance(result, UnquantizedLinearMethod)) - from vllm_ascend.quantization.methods import AscendW8A8DynamicLinearMethod, AscendW8A8DynamicFusedMoEMethod + from vllm_ascend.quantization.methods import AscendW8A8DynamicFusedMoEMethod, AscendW8A8DynamicLinearMethod + @patch("vllm_ascend.quantization.methods.AscendW8A8DynamicFusedMoEMethod.__init__") def test_get_moe_quant_method(self, mock_method): mock_method.return_value = None @@ -158,9 +157,7 @@ def test_get_moe_unquantized_method(self, mock_ignore_layer, mock_method): self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) self.assertTrue(isinstance(result, AscendUnquantizedFusedMoEMethod)) - def test_no_quant_method(self): layer = MagicMock(spec=Attention) result = self.config.get_quant_method(layer, "attn") self.assertIsNone(result) - diff --git a/tests/ut/quantization/test_method_adapters.py b/tests/ut/quantization/test_method_adapters.py index 05a6d90604e7..d7604df18423 100644 --- a/tests/ut/quantization/test_method_adapters.py +++ b/tests/ut/quantization/test_method_adapters.py @@ -1,20 +1,20 @@ from unittest.mock import MagicMock, patch import torch - -from tests.ut.base import TestBase from vllm.model_executor.layers.fused_moe import FusedMoeWeightScaleSupported from vllm.model_executor.layers.linear import ColumnParallelLinear + +from tests.ut.base import TestBase from vllm_ascend.quantization.method_adapters import ( - AscendLinearMethod, + AscendEmbeddingMethod, + AscendFusedMoEMethod, AscendKVCacheMethod, - AscendFusedMoEMethod, AscendEmbeddingMethod, + AscendLinearMethod, ) -from vllm_ascend.quantization.methods.base import AscendLinearScheme, AscendAttentionScheme, AscendMoEScheme +from vllm_ascend.quantization.methods.base import AscendAttentionScheme, AscendLinearScheme, AscendMoEScheme class TestAscendLinearMethod(TestBase): - def setUp(self): self.mock_scheme = MagicMock(spec=AscendLinearScheme) self.mock_scheme.get_weight.return_value = { @@ -105,7 +105,6 @@ def test_apply_delegates_to_scheme(self): class TestAscendKVCacheMethod(TestBase): - def setUp(self): self.mock_scheme = MagicMock(spec=AscendAttentionScheme) self.mock_scheme.create_weights.return_value = None @@ -129,15 +128,20 @@ def test_apply_delegates(self): value = torch.randn(4, 8, 64) self.mock_scheme.apply.return_value = torch.randn(4, 8, 64) self.method.apply( - layer, query, key, value, - kv_cache=None, attn_metadata=None, - attn_type=None, scale=1.0, output=None, + layer, + query, + key, + value, + kv_cache=None, + attn_metadata=None, + attn_type=None, + scale=1.0, + output=None, ) self.mock_scheme.apply.assert_called_once() class TestAscendFusedMoEMethod(TestBase): - def setUp(self): self.mock_scheme = MagicMock(spec=AscendMoEScheme) self.mock_scheme.group_size = 0 @@ -201,8 +205,6 @@ def test_create_weights_registers_parameters(self): self.assertEqual(layer.w13_weight_scale.quant_method, FusedMoeWeightScaleSupported.GROUP.value) self.assertEqual(layer.w2_weight_offset.quant_method, FusedMoeWeightScaleSupported.GROUP.value) - - def test_apply_method(self): layer = torch.nn.Module() x = torch.randn(8, 64) @@ -210,9 +212,7 @@ def test_apply_method(self): top_k = 3 renormalize = True self.mock_scheme.apply.return_value = None - self.method.apply( - layer, x, router_logits, top_k, renormalize - ) + self.method.apply(layer, x, router_logits, top_k, renormalize) self.mock_scheme.apply.assert_called_once() def test_supports_eplb_default_false(self): @@ -224,7 +224,6 @@ def test_supports_eplb_when_scheme_supports(self): class TestAscendEmbeddingMethod(TestBase): - def test_init(self): layer = MagicMock(spec=AscendLinearScheme) method = AscendEmbeddingMethod(layer) diff --git a/tests/ut/quantization/test_modelslim_config.py b/tests/ut/quantization/test_modelslim_config.py index 7cf8e4f1e4d4..a5d563a57ec9 100644 --- a/tests/ut/quantization/test_modelslim_config.py +++ b/tests/ut/quantization/test_modelslim_config.py @@ -287,7 +287,6 @@ def test_apply_extra_quant_adaptations_weight_packed(self): class TestGetPackedModulesMapping(TestBase): - def test_known_model_type(self): result = get_packed_modules_mapping("deepseek_v3") self.assertIsInstance(result, dict) @@ -300,7 +299,6 @@ def test_unknown_model_type(self): class TestGetLinearQuantType(TestBase): - def test_simple_non_fused_layer(self): quant_description = {"layer1.weight": "W8A8_DYNAMIC"} packed_modules_mapping = {} @@ -328,7 +326,6 @@ def test_fused_layer_inconsistent_quant_types(self): class TestGetQuantTypeForLayer(TestBase): - def test_attention_layer_with_fa_quant_type(self): quant_description = { "fa_quant_type": "C8", @@ -376,7 +373,6 @@ def test_with_none_packed_modules_mapping(self): class TestCreateSchemeForLayer(TestBase): - @patch("vllm_ascend.quantization.modelslim_config.get_scheme_class") def test_successful_scheme_creation(self, mock_get_scheme_class): mock_scheme_class = MagicMock() @@ -406,7 +402,6 @@ def test_unsupported_quant_type_raises_error(self, mock_get_scheme_class): class TestApplyVllmMapper(TestBase): - def test_apply_mapper_with_populated_quant_description(self): config = AscendModelSlimConfig({"old_key.weight": "INT8"}) mock_mapper = MagicMock() @@ -441,7 +436,6 @@ def test_empty_quant_description(self): class TestGetCacheScale(TestBase): - def test_c8_kv_cache_type_k_proj_scale(self): config = AscendModelSlimConfig({"kv_cache_type": "C8"}) result = config.get_cache_scale("model.layers.0.k_proj.kv_cache_scale") @@ -474,7 +468,6 @@ def test_no_matching_suffix(self): class TestQuantPrefixMapper(TestBase): - def test_sets_model_type_and_returns_prefix(self): config = AscendModelSlimConfig({}) result = config.quant_prefix_mapper("deepseek_v3", "model.layers.0") @@ -483,12 +476,13 @@ def test_sets_model_type_and_returns_prefix(self): class TestIsFaQuantLayer(TestBase): - def test_enable_fa_quant_true_matching_layer(self): - config = AscendModelSlimConfig({ - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - }) + config = AscendModelSlimConfig( + { + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + } + ) result = config.is_fa_quant_layer("layers.1.attn") self.assertTrue(result) @@ -498,29 +492,34 @@ def test_enable_fa_quant_false(self): self.assertFalse(result) def test_non_matching_layer_index(self): - config = AscendModelSlimConfig({ - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - }) + config = AscendModelSlimConfig( + { + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + } + ) result = config.is_fa_quant_layer("layers.2.attn") self.assertFalse(result) def test_no_layer_index(self): - config = AscendModelSlimConfig({ - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - }) + config = AscendModelSlimConfig( + { + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + } + ) result = config.is_fa_quant_layer("model.attn") self.assertFalse(result) class TestEnablingFaQuant(TestBase): - def test_decode_instance_with_fa_quant_layer(self): - config = AscendModelSlimConfig({ - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - }) + config = AscendModelSlimConfig( + { + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + } + ) mock_vllm_config = MagicMock() mock_vllm_config.kv_transfer_config.is_kv_consumer = True mock_vllm_config.kv_transfer_config.is_kv_producer = False @@ -529,10 +528,12 @@ def test_decode_instance_with_fa_quant_layer(self): self.assertTrue(result) def test_non_decode_instance(self): - config = AscendModelSlimConfig({ - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - }) + config = AscendModelSlimConfig( + { + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + } + ) mock_vllm_config = MagicMock() mock_vllm_config.kv_transfer_config.is_kv_consumer = False mock_vllm_config.kv_transfer_config.is_kv_producer = True @@ -541,10 +542,12 @@ def test_non_decode_instance(self): self.assertFalse(result) def test_is_fa_quant_layer_false(self): - config = AscendModelSlimConfig({ - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - }) + config = AscendModelSlimConfig( + { + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + } + ) mock_vllm_config = MagicMock() mock_vllm_config.kv_transfer_config.is_kv_consumer = True mock_vllm_config.kv_transfer_config.is_kv_producer = False @@ -554,12 +557,13 @@ def test_is_fa_quant_layer_false(self): class TestIsIndexerQuantLayer(TestBase): - def test_enable_indexer_quant_true_matching_layer(self): - config = AscendModelSlimConfig({ - "indexer_quant_type": "INT8", - "layers.1.indexer.quant_type": "INT8", - }) + config = AscendModelSlimConfig( + { + "indexer_quant_type": "INT8", + "layers.1.indexer.quant_type": "INT8", + } + ) result = config.is_indexer_quant_layer("layers.1.attn") self.assertTrue(result) @@ -569,29 +573,34 @@ def test_enable_indexer_quant_false(self): self.assertFalse(result) def test_non_matching_layer_index(self): - config = AscendModelSlimConfig({ - "indexer_quant_type": "INT8", - "layers.1.indexer.quant_type": "INT8", - }) + config = AscendModelSlimConfig( + { + "indexer_quant_type": "INT8", + "layers.1.indexer.quant_type": "INT8", + } + ) result = config.is_indexer_quant_layer("layers.2.attn") self.assertFalse(result) def test_no_layer_index(self): - config = AscendModelSlimConfig({ - "indexer_quant_type": "INT8", - "layers.1.indexer.quant_type": "INT8", - }) + config = AscendModelSlimConfig( + { + "indexer_quant_type": "INT8", + "layers.1.indexer.quant_type": "INT8", + } + ) result = config.is_indexer_quant_layer("model.attn") self.assertFalse(result) class TestGetKvQuantDtype(TestBase): - def test_enable_fa_quant_true_with_mla(self): - config = AscendModelSlimConfig({ - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - }) + config = AscendModelSlimConfig( + { + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + } + ) mock_model_config = MagicMock() mock_model_config.dtype = torch.float16 mock_model_config.use_mla = True @@ -601,10 +610,12 @@ def test_enable_fa_quant_true_with_mla(self): self.assertEqual(v_dtype, torch.float16) def test_enable_fa_quant_true_without_mla(self): - config = AscendModelSlimConfig({ - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - }) + config = AscendModelSlimConfig( + { + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + } + ) mock_model_config = MagicMock() mock_model_config.dtype = torch.float16 mock_model_config.use_mla = False @@ -625,14 +636,15 @@ def test_enable_fa_quant_false(self): class TestGetKvQuantSplitFactor(TestBase): - @patch("vllm_ascend.quantization.modelslim_config.calc_split_factor") def test_enable_fa_quant_true(self, mock_calc_split_factor): mock_calc_split_factor.return_value = 2.0 - config = AscendModelSlimConfig({ - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - }) + config = AscendModelSlimConfig( + { + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + } + ) kv_head_dim_list = [64, 64] result = config.get_kv_quant_split_factor("layers.1.attn", kv_head_dim_list) @@ -651,13 +663,14 @@ def test_enable_fa_quant_false(self, mock_calc_split_factor): class TestAddKvcacheQuantMetadata(TestBase): - def test_with_fa_quant_type(self): - config = AscendModelSlimConfig({ - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - "layers.2.fa_k.scale": "C8", - }) + config = AscendModelSlimConfig( + { + "fa_quant_type": "C8", + "layers.1.fa_k.scale": "C8", + "layers.2.fa_k.scale": "C8", + } + ) config._add_kvcache_quant_metadata() self.assertTrue(config.enable_fa_quant) @@ -667,11 +680,13 @@ def test_with_fa_quant_type(self): self.assertEqual(config.indexer_quant_layers, []) def test_with_indexer_quant_type(self): - config = AscendModelSlimConfig({ - "indexer_quant_type": "INT8", - "layers.1.indexer.quant_type": "INT8", - "layers.3.indexer.quant_type": "INT8", - }) + config = AscendModelSlimConfig( + { + "indexer_quant_type": "INT8", + "layers.1.indexer.quant_type": "INT8", + "layers.3.indexer.quant_type": "INT8", + } + ) config._add_kvcache_quant_metadata() self.assertFalse(config.enable_fa_quant) diff --git a/tests/ut/quantization/test_quant_parser.py b/tests/ut/quantization/test_quant_parser.py index 75e544c89f11..e6ca5de1f583 100644 --- a/tests/ut/quantization/test_quant_parser.py +++ b/tests/ut/quantization/test_quant_parser.py @@ -1,5 +1,7 @@ from unittest.mock import patch + import torch + from tests.ut.base import TestBase from vllm_ascend.quantization.quant_parser import ( QuantTypeMapping, @@ -10,7 +12,6 @@ class TestQuantTypeMapping(TestBase): - def test_get_quant_settings_returns_dict(self): settings = QuantTypeMapping.get_quant_settings() self.assertIsInstance(settings, dict) @@ -35,7 +36,6 @@ def test_w4a4_mxfp4_settings(self): class TestGetRollbackQuantType(TestBase): - def test_returns_down_proj_quant_type(self): config = { "model.layers.0.mlp.gate_proj": "W8A8_MXFP8", @@ -60,7 +60,6 @@ def test_returns_down_proj_type_with_multiple_entries(self): class TestParseMxfpQuantParams(TestBase): - def test_default_values(self): act, weight, scale, per_token, round_mode = parse_mxfp_quant_params() self.assertEqual(act, torch.float8_e4m3fn) @@ -89,35 +88,26 @@ def test_scale_type_overrides(self): class TestParseQuantMoeDownProjParams(TestBase): - @patch("vllm_ascend.quantization.quant_parser.ensure_mxfp8_scale_dtype_available") def test_w8a8_mxfp8_uses_rint_round_mode(self, mock_ensure): mock_ensure.return_value = None - act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params( - "W8A8_MXFP8", "round" - ) + act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params("W8A8_MXFP8", "round") self.assertEqual(round_mode, "rint") @patch("vllm_ascend.quantization.quant_parser.ensure_mxfp4_dtype_available") def test_w4a4_mxfp4_respects_parsed_round_mode(self, mock_ensure): mock_ensure.return_value = None - act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params( - "W4A4_MXFP4", "round" - ) + act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params("W4A4_MXFP4", "round") self.assertEqual(round_mode, "round") @patch("vllm_ascend.quantization.quant_parser.ensure_mxfp4_dtype_available") def test_w4a4_mxfp4_rint_round_mode(self, mock_ensure): mock_ensure.return_value = None - act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params( - "W4A4_MXFP4", "rint" - ) + act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params("W4A4_MXFP4", "rint") self.assertEqual(round_mode, "rint") @patch("vllm_ascend.quantization.quant_parser.ensure_mxfp8_scale_dtype_available") def test_w4a8_mxfp_uses_rint_round_mode(self, mock_ensure): mock_ensure.return_value = None - act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params( - "W4A8_MXFP", "round" - ) + act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params("W4A8_MXFP", "round") self.assertEqual(round_mode, "rint") diff --git a/tests/ut/quantization/test_quant_type.py b/tests/ut/quantization/test_quant_type.py index 820acd641f25..c70ab7298815 100644 --- a/tests/ut/quantization/test_quant_type.py +++ b/tests/ut/quantization/test_quant_type.py @@ -3,7 +3,6 @@ class TestQuantType(TestBase): - def test_enum_values(self): self.assertEqual(QuantType.NONE.value, 0) self.assertEqual(QuantType.W8A8.value, 1) diff --git a/tests/ut/quantization/test_utils.py b/tests/ut/quantization/test_utils.py index 526cb06c9614..0682e72c187b 100644 --- a/tests/ut/quantization/test_utils.py +++ b/tests/ut/quantization/test_utils.py @@ -161,7 +161,6 @@ def test_passes_revision_to_detect(self, mock_detect, mock_get_quant): class TestEnableFaQuant(TestBase): - def test_non_quantization_scenarios(self): vllm_config = MagicMock() vllm_config.quant_config = None From 9020d21cb5495f8612b23430cd58d44fa08f8c27 Mon Sep 17 00:00:00 2001 From: Wang Kunpeng <1289706727@qq.com> Date: Fri, 24 Apr 2026 22:11:35 +0800 Subject: [PATCH 3/4] [Test]Add quantization test case Signed-off-by: Wang Kunpeng <1289706727@qq.com> --- tests/ut/quantization/methods/test_kv_c8.py | 22 -- .../ut/quantization/methods/test_registry.py | 28 +- tests/ut/quantization/methods/test_w4a16.py | 39 --- .../methods/test_w4a4_flatquant.py | 4 +- .../methods/test_w4a4_laos_dynamic.py | 37 --- .../quantization/methods/test_w4a4_mxfp4.py | 27 +- tests/ut/quantization/methods/test_w4a8.py | 6 - tests/ut/quantization/methods/test_w8a16.py | 41 --- .../quantization/methods/test_w8a8_dynamic.py | 187 +---------- .../quantization/methods/test_w8a8_mxfp8.py | 88 +---- .../quantization/methods/test_w8a8_pdmix.py | 39 +-- .../quantization/methods/test_w8a8_static.py | 25 -- .../test_compressed_tensors_config.py | 48 +-- tests/ut/quantization/test_method_adapters.py | 55 +--- .../ut/quantization/test_modelslim_config.py | 309 +----------------- tests/ut/quantization/test_quant_parser.py | 74 +---- tests/ut/quantization/test_quant_type.py | 28 -- tests/ut/quantization/test_utils.py | 7 +- vllm_ascend/quantization/method_adapters.py | 1 - vllm_ascend/quantization/methods/kv_c8.py | 2 - vllm_ascend/quantization/methods/w4a16.py | 58 ++-- .../quantization/methods/w4a4_laos_dynamic.py | 5 +- 22 files changed, 80 insertions(+), 1050 deletions(-) delete mode 100644 tests/ut/quantization/test_quant_type.py diff --git a/tests/ut/quantization/methods/test_kv_c8.py b/tests/ut/quantization/methods/test_kv_c8.py index 46f8a14875f5..1b350df27ef4 100644 --- a/tests/ut/quantization/methods/test_kv_c8.py +++ b/tests/ut/quantization/methods/test_kv_c8.py @@ -166,37 +166,15 @@ def tearDown(self): def test_init_with_full_config(self): """Test initialization when config has all attributes""" method = self.method_class() - - self.assertTrue(method.transpose_weight) - self.assertFalse(method.printFlag) self.assertEqual(method.kv_lora_rank, 128) self.assertEqual(method.qk_rope_head_dim, 64) - def test_init_without_kv_lora_rank(self): - """Test initialization when config lacks kv_lora_rank""" - delattr(self.mock_hf_config, "kv_lora_rank") - - method = self.method_class() - - self.assertEqual(method.kv_lora_rank, 0) - self.assertEqual(method.qk_rope_head_dim, 64) - - def test_init_without_qk_rope_head_dim(self): - """Test initialization when config lacks qk_rope_head_dim""" - delattr(self.mock_hf_config, "qk_rope_head_dim") - - method = self.method_class() - - self.assertEqual(method.kv_lora_rank, 128) - self.assertEqual(method.qk_rope_head_dim, 0) - def test_init_without_both_attributes(self): """Test initialization when config lacks both attributes""" delattr(self.mock_hf_config, "kv_lora_rank") delattr(self.mock_hf_config, "qk_rope_head_dim") method = self.method_class() - self.assertEqual(method.kv_lora_rank, 0) self.assertEqual(method.qk_rope_head_dim, 0) diff --git a/tests/ut/quantization/methods/test_registry.py b/tests/ut/quantization/methods/test_registry.py index 5aa145bf2636..6ea88d406be1 100644 --- a/tests/ut/quantization/methods/test_registry.py +++ b/tests/ut/quantization/methods/test_registry.py @@ -23,37 +23,33 @@ def apply(self, layer, x, bias=None, tp_rank=0): scheme_class = get_scheme_class("TEST_QUANT_TYPE", "linear") self.assertIs(scheme_class, TestLinearScheme) + def test_register_scheme_duplicate_raises(self): + with self.assertRaises(ValueError): + + @register_scheme("W8A8_DYNAMIC", "linear") + class Duplicate: + pass + class TestGetSchemeClass(TestBase): - def test_get_scheme_class_existing_linear(self): + def test_get_existing_scheme_class_existing_linear(self): cls = get_scheme_class("W8A8_DYNAMIC", "linear") self.assertIsNotNone(cls) self.assertTrue(issubclass(cls, AscendLinearScheme)) - def test_get_scheme_class_existing_moe(self): cls = get_scheme_class("W8A8_DYNAMIC", "moe") self.assertIsNotNone(cls) self.assertTrue(issubclass(cls, AscendMoEScheme)) - def test_get_scheme_class_existing_attention(self): cls = get_scheme_class("FAKQuant", "attention") self.assertIsNotNone(cls) - def test_get_scheme_class_nonexistent(self): + def test_get_nonexistent_scheme_class(self): cls = get_scheme_class("NONEXISTENT", "linear") self.assertIsNone(cls) - - def test_get_scheme_class_nonexistent_layer_type(self): cls = get_scheme_class("W8A8_DYNAMIC", "nonexistent") self.assertIsNone(cls) - def test_register_scheme_duplicate_raises(self): - with self.assertRaises(ValueError): - - @register_scheme("W8A8_DYNAMIC", "linear") - class Duplicate: - pass - def test_all_linear_schemes_subclass_ascend_linear_scheme(self): for (quant_type, layer_type), scheme_cls in _SCHEME_REGISTRY.items(): if layer_type == "linear": @@ -72,9 +68,3 @@ def test_all_moe_schemes_subclass_ascend_moe_scheme(self): def test_registry_not_empty(self): self.assertGreater(len(_SCHEME_REGISTRY), 0) - - def test_registry_key_format(self): - for key in _SCHEME_REGISTRY.keys(): - self.assertEqual(len(key), 2) - self.assertIsInstance(key[0], str) - self.assertIsInstance(key[1], str) diff --git a/tests/ut/quantization/methods/test_w4a16.py b/tests/ut/quantization/methods/test_w4a16.py index 768dd8930555..2fcdb90f69f0 100644 --- a/tests/ut/quantization/methods/test_w4a16.py +++ b/tests/ut/quantization/methods/test_w4a16.py @@ -107,13 +107,6 @@ def setUp(self, mock_get_current_vllm_config, mock_get_ascend_config): self.quant_method = AscendW4A16FusedMoEMethod() - def test_init(self): - self.assertTrue(self.quant_method.transpose_weight) - self.assertEqual(self.quant_method.num_bits, 4) - self.assertEqual(self.quant_method.pack_factor, 8) - self.assertEqual(self.quant_method.group_size, self.group_size) - self.assertFalse(self.quant_method.dynamic_eplb) - def test_get_weight(self): param_dict = self.quant_method.get_weight(self.experts, self.input_size, self.output_size, torch.bfloat16) @@ -134,22 +127,12 @@ def test_get_dynamic_quant_param(self): expected_w13_scale_shape = (self.experts, 2 * self.input_size, self.output_size // self.group_size) self.assertEqual(param_dict["w13_weight_scale"].shape, expected_w13_scale_shape) - self.assertEqual(param_dict["w2_weight_scale"].dtype, torch.bfloat16) - expected_w2_scale_shape = (self.experts, self.output_size, self.input_size // self.group_size) - self.assertEqual(param_dict["w2_weight_scale"].shape, expected_w2_scale_shape) - - self.assertEqual(param_dict["w13_weight_shape"].dtype, torch.int32) - self.assertEqual(param_dict["w13_weight_shape"].shape, (self.experts, 2)) - self.assertEqual(param_dict["w2_weight_shape"].dtype, torch.int32) self.assertEqual(param_dict["w2_weight_shape"].shape, (self.experts, 2)) self.assertEqual(param_dict["w13_weight_offset"].dtype, torch.bfloat16) self.assertEqual(param_dict["w13_weight_offset"].shape, expected_w13_scale_shape) - self.assertEqual(param_dict["w2_weight_offset"].dtype, torch.bfloat16) - self.assertEqual(param_dict["w2_weight_offset"].shape, expected_w2_scale_shape) - def build_layer(self): """Build a mock layer for testing""" layer = torch.nn.Module() @@ -199,36 +182,14 @@ def mock_convert_weight(weight): return torch.zeros(new_shape, dtype=torch.int32) mock_npu_convert_weight_to_int4pack.side_effect = mock_convert_weight - layer = self.build_layer() - self.quant_method.transpose_weight = True - self.quant_method.process_weights_after_loading(layer) self.assertEqual(layer.w13_weight_packed.data.shape, torch.Size([8, 128, 8])) self.assertEqual(layer.w2_weight_packed.data.shape, torch.Size([8, 32, 16])) - self.assertEqual(layer.w13_weight_scale.data.shape, torch.Size([8, 4, 64])) - self.assertEqual(layer.w2_weight_scale.data.shape, torch.Size([8, 1, 128])) - self.assertEqual(layer.w13_weight_offset.data.shape, torch.Size([8, 4, 64])) self.assertEqual(layer.w2_weight_offset.data.shape, torch.Size([8, 1, 128])) - self.assertTrue(layer.w13_weight_scale.data.is_contiguous()) - self.assertTrue(layer.w2_weight_scale.data.is_contiguous()) - self.assertTrue(layer.w13_weight_offset.data.is_contiguous()) - self.assertTrue(layer.w2_weight_offset.data.is_contiguous()) - - def test_process_weights_after_loading_without_transpose(self): - layer = self.build_layer() - self.quant_method.transpose_weight = False - - original_w13_data = layer.w13_weight_packed.data.clone() - original_w2_data = layer.w2_weight_packed.data.clone() - - self.quant_method.process_weights_after_loading(layer) - - self.assertTrue(torch.equal(layer.w13_weight_packed.data, original_w13_data)) - self.assertTrue(torch.equal(layer.w2_weight_packed.data, original_w2_data)) @patch("vllm_ascend.quantization.methods.w4a16._EXTRA_CTX") @patch("vllm_ascend.quantization.methods.w4a16.select_experts") diff --git a/tests/ut/quantization/methods/test_w4a4_flatquant.py b/tests/ut/quantization/methods/test_w4a4_flatquant.py index f48b250d52d1..c00c791be49c 100644 --- a/tests/ut/quantization/methods/test_w4a4_flatquant.py +++ b/tests/ut/quantization/methods/test_w4a4_flatquant.py @@ -72,7 +72,7 @@ def test_large_batch_multiple_calls(self, mock_npu): (torch.randint(0, 255, (last_chunk_size, 24, 4), dtype=torch.int32), torch.randn(last_chunk_size)) ) mock_npu.npu_kronecker_quant.side_effect = mock_returns - result_x, result_scale = batched_kronecker_quant(x, left_trans, right_trans, 0.95) + batched_kronecker_quant(x, left_trans, right_trans, 0.95) self.assertEqual(mock_npu.npu_kronecker_quant.call_count, num_chunks) @patch("vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu") @@ -85,7 +85,7 @@ def test_exact_max_batch_size(self, mock_npu): torch.randint(0, 255, (batch_size, 24, 4), dtype=torch.int32), torch.randn(batch_size, dtype=torch.float32), ) - result_x, result_scale = batched_kronecker_quant(x, left_trans, right_trans, 0.95) + batched_kronecker_quant(x, left_trans, right_trans, 0.95) mock_npu.npu_kronecker_quant.assert_called_once() ## Test AscendW4A4FlatQuantDynamicLinearMethod Class diff --git a/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py b/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py index 584db5fdbdbb..2dc7045953f2 100644 --- a/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py +++ b/tests/ut/quantization/methods/test_w4a4_laos_dynamic.py @@ -11,9 +11,6 @@ class TestAscendW4A4LaosDynamicLinearMethod(TestBase): def setUp(self): self.method = AscendW4A4LaosDynamicLinearMethod() - def test_init_transpose_weight_true(self): - self.assertTrue(self.method.transpose_weight) - def test_get_weight_various_sizes(self): sizes = [(64, 128), (256, 512), (1024, 2048)] for input_size, output_size in sizes: @@ -30,30 +27,6 @@ def test_get_perchannel_param_various_output_sizes(self): self.assertEqual(result["weight_scale"].dtype, torch.float32) self.assertEqual(result["weight_offset"].dtype, torch.float32) - def test_get_pertensor_param_empty(self): - result = self.method.get_pertensor_param(torch.bfloat16) - self.assertEqual(result, {}) - - def test_get_pergroup_param_various_layer_types(self): - for layer_type in ["column", "row", None]: - result = self.method.get_pergroup_param(256, 128, torch.bfloat16, layer_type) - self.assertEqual(result, {}) - - @patch("torch_npu.npu_quant_matmul") - @patch("torch_npu.npu_dynamic_quant") - def test_apply_2d_input(self, mock_dyn_quant, mock_matmul): - mock_dyn_quant.return_value = ( - torch.randint(0, 15, (32, 128), dtype=torch.int32), - torch.randn(32, dtype=torch.float32), - ) - mock_matmul.return_value = torch.randn(32, 256) - layer = MagicMock() - layer.weight = MagicMock(data=torch.randint(-8, 7, (256, 128), dtype=torch.int8)) - layer.weight_scale = MagicMock(data=torch.randn(256, dtype=torch.float32)) - x = torch.randn(32, 128, dtype=torch.bfloat16) - output = self.method.apply(layer, x) - mock_matmul.assert_called_once() - @patch("torch_npu.npu_quant_matmul") @patch("torch_npu.npu_dynamic_quant") def test_apply_with_bias(self, mock_dyn_quant, mock_matmul): @@ -72,16 +45,6 @@ def test_apply_with_bias(self, mock_dyn_quant, mock_matmul): expected_output = expected_output + bias self.assertTrue(torch.equal(output, expected_output)) - @patch("torch_npu.npu_convert_weight_to_int4pack") - def test_process_weights_no_transpose_when_false(self, mock_convert): - mock_convert.return_value = torch.randint(0, 15, (128, 32), dtype=torch.int32) - self.method.transpose_weight = False - layer = nn.Module() - layer.weight = nn.Parameter(torch.randint(-8, 7, (128, 256), dtype=torch.int8), requires_grad=False) - layer.weight_scale = nn.Parameter(torch.randn(128, 1, dtype=torch.float32), requires_grad=False) - self.method.process_weights_after_loading(layer) - self.assertEqual(layer.weight.shape, (128, 32)) - @patch("torch_npu.npu_convert_weight_to_int4pack") def test_process_weights_various_input_sizes(self, mock_convert): for input_size, output_size in [(64, 128), (256, 512)]: diff --git a/tests/ut/quantization/methods/test_w4a4_mxfp4.py b/tests/ut/quantization/methods/test_w4a4_mxfp4.py index d594cd04c5cb..a703a855df38 100644 --- a/tests/ut/quantization/methods/test_w4a4_mxfp4.py +++ b/tests/ut/quantization/methods/test_w4a4_mxfp4.py @@ -5,7 +5,6 @@ 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.base import QuantType from vllm_ascend.quantization.methods.w4a4_mxfp4 import ( AscendW4A4MXFP4DynamicFusedMoEMethod, AscendW4A4MXFP4DynamicLinearMethod, @@ -20,9 +19,6 @@ def setUp(self, mock_vllm, mock_ensure): mock_ensure.return_value = None self.scheme = AscendW4A4MXFP4DynamicLinearMethod() - def test_model_dtype_is_none(self): - self.assertIsNone(self.scheme.model_dtype) - 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) @@ -37,14 +33,6 @@ def test_get_pergroup_param_based_on_group_size(self): self.assertEqual(result["weight_scale"].shape, (128, 256 // gs)) self.assertEqual(result["weight_scale"].dtype, torch.uint8) - def test_get_pertensor_param_returns_empty(self): - result = self.scheme.get_pertensor_param(torch.bfloat16) - self.assertEqual(result, {}) - - def test_get_perchannel_param_returns_empty(self): - result = self.scheme.get_perchannel_param(128, torch.bfloat16) - self.assertEqual(result, {}) - 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) @@ -53,13 +41,6 @@ def test_process_weights_after_loading_transposes(self): self.assertEqual(layer.weight.shape, (128, 128)) self.assertEqual(layer.weight_scale.shape[0], 4) - @patch("vllm_ascend.quantization.methods.w4a4_mxfp4.ensure_mxfp4_linear_available") - @patch("vllm_ascend.quantization.methods.w4a4_mxfp4.get_current_vllm_config") - def test_group_size_from_config(self, mock_vllm_config, mock_ensure): - mock_vllm_config.return_value = create_mock_vllm_config(quant_description={"group_size": 64}) - scheme = AscendW4A4MXFP4DynamicLinearMethod() - self.assertEqual(scheme.group_size, 64) - @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 = ( @@ -92,10 +73,6 @@ def setUp(self, mock_ascend, mock_vllm, mock_ep, mock_ensure): mock_ensure.return_value = None self.scheme = AscendW4A4MXFP4DynamicFusedMoEMethod() - def test_init(self): - self.assertEqual(self.scheme.quant_type, QuantType.MXFP4) - self.assertIsNone(self.scheme.model_dtype) - 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) @@ -126,9 +103,7 @@ def test_process_weights_transposes_weights(self): 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.w2_weight.shape, (8, 128, 128)) self.assertEqual(layer.w13_weight_scale.shape, (8, 2, 256, 2)) - self.assertEqual(layer.w2_weight_scale.shape, (8, 4, 128, 2)) @patch("vllm_ascend.quantization.methods.w4a4_mxfp4.torch_npu") @patch("vllm_ascend.quantization.methods.w4a4_mxfp4._EXTRA_CTX") @@ -153,7 +128,7 @@ def test_apply_full_params(self, mock_select, mock_ctx, mock_npu): 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() - result = self.scheme.apply( + self.scheme.apply( layer, x, router_logits, diff --git a/tests/ut/quantization/methods/test_w4a8.py b/tests/ut/quantization/methods/test_w4a8.py index 374f4b2b954c..c80cdd5a8c77 100644 --- a/tests/ut/quantization/methods/test_w4a8.py +++ b/tests/ut/quantization/methods/test_w4a8.py @@ -290,9 +290,6 @@ def test_process_weights_after_loading_compressed_tensors(self, mock_npu, mock_n 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) @patch("vllm_ascend.quantization.methods.w4a8._EXTRA_CTX") @patch("vllm_ascend.quantization.methods.w4a8.select_experts") @@ -337,9 +334,6 @@ def test_apply_comprehensive(self, mock_build_input, mock_select, mock_ctx): use_grouped_topk=False, global_num_experts=num_experts, expert_map=expert_map, - topk_group=None, - num_expert_group=None, - custom_routing_function=None, scoring_func="softmax", routed_scaling_factor=1.0, e_score_correction_bias=e_score_correction_bias, diff --git a/tests/ut/quantization/methods/test_w8a16.py b/tests/ut/quantization/methods/test_w8a16.py index 755de13e4dc3..ab4f5e0c4416 100644 --- a/tests/ut/quantization/methods/test_w8a16.py +++ b/tests/ut/quantization/methods/test_w8a16.py @@ -32,14 +32,6 @@ def test_get_per_channel_param(self): self.assertEqual(per_channel_params["weight_offset"].shape, (output_size, 1)) self.assertEqual(len(per_channel_params), 2) - def test_get_pertensor_param_returns_empty(self): - result = self.method.get_pertensor_param(torch.bfloat16) - self.assertEqual(result, {}) - - def test_get_pergroup_param_returns_empty(self): - result = self.method.get_pergroup_param(256, 128, torch.bfloat16) - self.assertEqual(result, {}) - @patch("torch_npu.npu_weight_quant_batchmatmul") def test_apply_with_x_is_int8(self, mock_npu_weight_quant_batchmatmul): layer = MagicMock() @@ -57,22 +49,6 @@ def test_apply_with_x_is_int8(self, mock_npu_weight_quant_batchmatmul): self.assertTrue(torch.equal(output, expected_y_output)) mock_npu_weight_quant_batchmatmul.assert_called_once() - @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "0"}) - @patch("torch_npu.npu_format_cast") - def test_process_weights_after_loading_with_nz0(self, mock_npu_format_cast): - layer = MagicMock() - layer.weight.data = torch.randint(-128, 127, (128, 256), dtype=torch.int8) - layer.weight_scale.data = torch.randn(128, 1) - layer.weight_offset.data = torch.randn(128, 1) - - mock_npu_format_cast.side_effect = identity - self.method.process_weights_after_loading(layer) - - self.assertEqual(layer.weight.data.shape, (256, 128)) - self.assertEqual(layer.weight_scale.data.shape, (128,)) - self.assertEqual(layer.weight_offset.data.shape, (128,)) - mock_npu_format_cast.assert_not_called() - @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "1"}) @patch("torch_npu.npu_format_cast") def test_process_weights_after_loading_with_nz1(self, mock_npu_format_cast): @@ -89,20 +65,3 @@ def test_process_weights_after_loading_with_nz1(self, mock_npu_format_cast): self.assertEqual(layer.weight_scale.data.shape, (128,)) self.assertEqual(layer.weight_offset.data.shape, (128,)) mock_npu_format_cast.assert_called_once() - - @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "2"}) - @patch("torch_npu.npu_format_cast") - def test_process_weights_after_loading_with_nz2(self, mock_npu_format_cast): - layer = MagicMock() - - layer.weight.data = torch.randint(-128, 127, (128, 256), dtype=torch.int8) - layer.weight_scale.data = torch.randn(128, 1) - layer.weight_offset.data = torch.randn(128, 1) - - mock_npu_format_cast.side_effect = identity - self.method.process_weights_after_loading(layer) - - self.assertEqual(layer.weight.data.shape, (256, 128)) - self.assertEqual(layer.weight_scale.data.shape, (128,)) - self.assertEqual(layer.weight_offset.data.shape, (128,)) - mock_npu_format_cast.assert_called_once() diff --git a/tests/ut/quantization/methods/test_w8a8_dynamic.py b/tests/ut/quantization/methods/test_w8a8_dynamic.py index e84e2bb1a503..e9998b2eaabf 100644 --- a/tests/ut/quantization/methods/test_w8a8_dynamic.py +++ b/tests/ut/quantization/methods/test_w8a8_dynamic.py @@ -12,23 +12,7 @@ from vllm_ascend.quantization.methods.w8a8_dynamic import ( AscendW8A8DynamicFusedMoEMethod, AscendW8A8DynamicLinearMethod, - scale_from_float_to_int64, ) -from vllm_ascend.quantization.quant_type import QuantType - - -class TestScaleFromFloatToInt64(TestBase): - def test_scale_conversion_basic(self): - scale = torch.tensor([0.5], dtype=torch.float32) - with patch.object(scale, "cpu", return_value=scale), patch.object(scale, "to", return_value=scale): - result = scale_from_float_to_int64(scale) - self.assertEqual(result.dtype, torch.int64) - - def test_scale_conversion_preserves_device(self): - scale = torch.tensor([0.5, 1.0], dtype=torch.float32) - with patch.object(scale, "cpu", return_value=scale), patch.object(scale, "to", return_value=scale): - result = scale_from_float_to_int64(scale) - self.assertEqual(result.shape, (2,)) class TestAscendW8A8DynamicLinearMethod(TestBase): @@ -51,22 +35,6 @@ def test_get_perchannel_param_dtype_variations(self): self.assertEqual(params["weight_scale"].shape, (128, 1)) self.assertEqual(params["weight_offset"].shape, (128, 1)) - @patch("torch_npu.npu_quant_matmul") - @patch("torch_npu.npu_dynamic_quant") - def test_apply_2d_input(self, mock_dyn_quant, mock_matmul): - mock_dyn_quant.return_value = ( - torch.randint(-128, 127, (32, 128), dtype=torch.int8), - torch.randn(32, dtype=torch.float32), - ) - mock_matmul.return_value = torch.randn(32, 256) - layer = MagicMock() - layer.weight = torch.randint(-128, 127, (128, 256), dtype=torch.int8) - layer.weight_scale = torch.randn(256, dtype=torch.float32) - x = torch.randn(32, 128, dtype=torch.bfloat16) - output = self.method.apply(layer, x) - mock_dyn_quant.assert_called_once() - mock_matmul.assert_called_once() - @patch("torch_npu.npu_quant_matmul") @patch("torch_npu.npu_dynamic_quant") def test_apply_3d_input_with_squeeze(self, mock_dyn_quant, mock_matmul): @@ -80,6 +48,8 @@ def test_apply_3d_input_with_squeeze(self, mock_dyn_quant, mock_matmul): layer.weight_scale = torch.randn(256, dtype=torch.float32) x = torch.randn(32, 1, 128, dtype=torch.bfloat16) output = self.method.apply(layer, x) + mock_dyn_quant.assert_called_once() + mock_matmul.assert_called_once() self.assertEqual(output.shape, (32, 1, 1, 256)) def test_process_weights_after_loading(self): @@ -117,9 +87,6 @@ def setUp(self, mock_ep, mock_ascend, mock_mc2, mock_rank): mock_rank.return_value = 0 self.quant_method = AscendW8A8DynamicFusedMoEMethod() - def test_quant_type_is_w8a8(self): - self.assertEqual(self.quant_method.quant_type, QuantType.W8A8) - def test_get_weight_various_expert_counts(self): expert_counts = [4, 8, 16, 32] for num_experts in expert_counts: @@ -195,156 +162,10 @@ def test_apply_uses_explicit_dispatch_and_mlp_args(self, mock_select_experts, mo self.assertIs(fused_experts_input.topk_weights, topk_weights) self.assertIs(fused_experts_input.topk_ids, topk_ids) - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_flash_common3_context") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts") - def test_apply_overlap_gate_uses_fc3_context( - self, - mock_select_experts, - mock_extra_ctx, - mock_get_flash_common3_context, - ): - tokens = 4 - hidden_size = self.hidden_size - layer = torch.nn.Module() - layer.w13_weight = torch.randint( - -8, - 8, - (self.num_experts, 2 * self.intermediate_size, hidden_size), - dtype=torch.int8, - ) - layer.w2_weight = torch.randint( - -8, - 8, - (self.num_experts, hidden_size, self.intermediate_size), - dtype=torch.int8, - ) - layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size, dtype=torch.float32) - layer.w2_weight_scale = torch.ones(self.num_experts, hidden_size, dtype=torch.float32) - - x = torch.randn(tokens, hidden_size, dtype=torch.float32) - router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32) - topk_weights = torch.randn(tokens, 2, dtype=torch.float32) - topk_ids = torch.randint(0, self.num_experts, (tokens, 2), dtype=torch.int64) - mc2_mask = torch.tensor([1, 0, 1, 0], dtype=torch.bool) - pertoken_scale = torch.randn(tokens, dtype=torch.float32) - - self.quant_method.multistream_overlap_gate = True - self.quant_method.in_dtype = torch.float32 - mock_get_flash_common3_context.return_value = Mock(topk_weights=topk_weights, topk_ids=topk_ids) - - mock_comm = Mock() - mock_comm.fused_experts.return_value = torch.randn(tokens, hidden_size, dtype=torch.float32) - mock_extra_ctx.moe_comm_method = mock_comm - mock_extra_ctx.moe_comm_type = MoECommType.ALLGATHER - - self.quant_method.apply( - layer=layer, - x=x, - router_logits=router_logits, - top_k=2, - renormalize=True, - global_num_experts=self.num_experts, - activation="gelu", - apply_router_weight_on_input=True, - mc2_mask=mc2_mask, - pertoken_scale=pertoken_scale, - ) - - mock_select_experts.assert_not_called() - fused_experts_input = mock_comm.fused_experts.call_args.kwargs["fused_experts_input"] - self.assertEqual(fused_experts_input.activation, "gelu") - self.assertTrue(fused_experts_input.routing.apply_router_weight_on_input) - self.assertIs(fused_experts_input.routing.mc2_mask, mc2_mask) - self.assertIs(fused_experts_input.routing.pertoken_scale, pertoken_scale) - self.assertIs(fused_experts_input.topk_weights, topk_weights) - self.assertIs(fused_experts_input.topk_ids, topk_ids) - - @patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.zero_experts_compute") - def test_apply_with_zero_experts(self, mock_zero, mock_select, mock_ctx): - tokens = 4 - layer = MagicMock() - layer.w13_weight = torch.randint( - -8, 8, (self.num_experts, 2 * self.intermediate_size, self.hidden_size), dtype=torch.int8 - ) - layer.w2_weight = torch.randint( - -8, 8, (self.num_experts, self.hidden_size, self.intermediate_size), dtype=torch.int8 - ) - layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size) - layer.w2_weight_scale = torch.ones(self.num_experts, self.hidden_size) - layer.zero_expert_num = 2 - layer.zero_expert_type = "shared" - layer.n_shared_experts = 0 - layer.mix_placement = False - x = torch.randn(tokens, self.hidden_size, dtype=torch.float32) - 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_zero.return_value = (topk_ids, topk_weights, torch.randn(tokens, self.hidden_size)) - 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.quant_method.in_dtype = torch.float32 - self.quant_method.apply(layer, x, router_logits, top_k=2, renormalize=True, global_num_experts=self.num_experts) - - @patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts") - def test_apply_with_enable_force_load_balance(self, mock_select, mock_ctx): - tokens = 4 - layer = MagicMock() - layer.w13_weight = torch.randint( - -8, 8, (self.num_experts, 2 * self.intermediate_size, self.hidden_size), dtype=torch.int8 - ) - layer.w2_weight = torch.randint( - -8, 8, (self.num_experts, self.hidden_size, self.intermediate_size), dtype=torch.int8 - ) - layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size) - layer.w2_weight_scale = torch.ones(self.num_experts, self.hidden_size) - layer.zero_expert_num = 0 - layer.zero_expert_type = None - layer.n_shared_experts = 0 - layer.mix_placement = False - x = torch.randn(tokens, self.hidden_size, dtype=torch.float32) - 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.quant_method.in_dtype = torch.float32 - self.quant_method.apply( - layer, - x, - router_logits, - top_k=2, - renormalize=True, - global_num_experts=self.num_experts, - enable_force_load_balance=True, - ) - @patch("torch_npu.npu_format_cast") @patch("vllm_ascend.quantization.methods.w8a8_dynamic.envs_ascend") - def test_process_weights_with_fused_mc2(self, mock_envs, mock_format_cast): + def test_process_weights_after_loading(self, mock_envs, mock_format_cast): mock_envs.VLLM_ASCEND_ENABLE_FUSED_MC2 = 1 - mock_format_cast.return_value = torch.randint( - -8, 8, (self.num_experts, self.hidden_size, 2 * self.intermediate_size), dtype=torch.int8 - ) - layer = create_moe_layer( - num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size - ) - self.quant_method.process_weights_after_loading(layer) - self.assertTrue(hasattr(layer, "w13_weight_scale_fp32")) - - @patch("torch_npu.npu_format_cast") - @patch("vllm_ascend.quantization.methods.w8a8_dynamic.envs_ascend") - def test_process_weights_with_dynamic_eplb(self, mock_envs, mock_format_cast): - mock_envs.VLLM_ASCEND_ENABLE_FUSED_MC2 = 0 self.quant_method.dynamic_eplb = True mock_format_cast.return_value = torch.randint( -8, 8, (self.num_experts, self.hidden_size, 2 * self.intermediate_size), dtype=torch.int8 @@ -353,4 +174,6 @@ def test_process_weights_with_dynamic_eplb(self, mock_envs, mock_format_cast): num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size ) self.quant_method.process_weights_after_loading(layer) + self.assertEqual(layer.fused_w1_scale.dtype, torch.int64) self.assertTrue(hasattr(layer, "w13_weight_list")) + self.assertFalse(hasattr(layer, "w13_weight_scale_fp32")) diff --git a/tests/ut/quantization/methods/test_w8a8_mxfp8.py b/tests/ut/quantization/methods/test_w8a8_mxfp8.py index 81188414b612..f46d17bda8bd 100644 --- a/tests/ut/quantization/methods/test_w8a8_mxfp8.py +++ b/tests/ut/quantization/methods/test_w8a8_mxfp8.py @@ -9,7 +9,6 @@ create_mock_vllm_config, create_mxfp_moe_layer, ) -from vllm_ascend.quantization.methods.base import QuantType from vllm_ascend.quantization.methods.w8a8_mxfp8 import ( AscendW8A8MXFP8DynamicFusedMoEMethod, AscendW8A8MXFP8DynamicLinearMethod, @@ -24,20 +23,12 @@ def setUp(self, mock_vllm, mock_ensure): mock_ensure.return_value = None self.scheme = AscendW8A8MXFP8DynamicLinearMethod() - def test_init(self): - self.assertIsNone(self.scheme.model_dtype) - self.assertEqual(self.scheme.group_size, 32) - - def test_get_weight_float8_dtype(self): - result = self.scheme.get_weight(256, 128, torch.bfloat16) - self.assertEqual(result["weight"].dtype, torch.float8_e4m3fn) - self.assertEqual(len(result), 1) - def test_get_weight_various_input_sizes(self): sizes = [(128, 64), (512, 256), (1024, 512)] for input_size, output_size in sizes: result = self.scheme.get_weight(input_size, output_size, torch.bfloat16) self.assertEqual(result["weight"].shape, (output_size, input_size)) + self.assertEqual(result["weight"].dtype, torch.float8_e4m3fn) def test_get_pergroup_param_group_size_variations(self): group_sizes = [16, 32, 64, 128] @@ -55,31 +46,7 @@ def test_process_weights_stores_original_shapes(self): self.assertTrue(hasattr(layer, "_mxfp8_original_shapes")) self.assertEqual(layer._mxfp8_original_shapes["weight"], (128, 256)) self.assertTrue(layer._mxfp8_transformed) - - def test_process_weights_double_call_no_change(self): - layer = nn.Module() - layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), 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) - weight_after_first = layer.weight.data.clone() - self.scheme.process_weights_after_loading(layer) - self.assertTrue(torch.equal(layer.weight.data, weight_after_first)) - - def test_restore_without_mxfp8_transformed(self): - layer = nn.Module() - layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), requires_grad=False) - original_weight_shape = layer.weight.shape - layer._mxfp8_transformed = False - self.scheme.restore_weights_for_rl_loading(layer) - self.assertEqual(layer.weight.shape, original_weight_shape) - - def test_restore_without_original_shapes_raises(self): - layer = nn.Module() - layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), requires_grad=False) - layer.weight_scale = nn.Parameter(torch.randint(0, 255, (128, 8), dtype=torch.uint8), requires_grad=False) - layer._mxfp8_transformed = True - with self.assertRaises(RuntimeError): - self.scheme.restore_weights_for_rl_loading(layer) + self.assertEqual(layer.weight_scale.shape, (4, 128, 2)) def test_restore_after_process_returns_original_shape(self): layer = nn.Module() @@ -93,13 +60,6 @@ def test_restore_after_process_returns_original_shape(self): self.assertEqual(layer.weight_scale.shape, original_scale_shape) self.assertFalse(layer._mxfp8_transformed) - def test_scale_shape_transformation(self): - layer = nn.Module() - layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), 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_scale.shape, (4, 128, 2)) - @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.torch_npu") def test_apply(self, mock_torch_npu): from vllm_ascend.device.mxfp_compat import FLOAT8_E8M0FNU_DTYPE @@ -121,7 +81,6 @@ def test_apply(self, mock_torch_npu): self.assertEqual(call_kwargs["bias"].dtype, torch.float32) self.assertEqual(call_kwargs["group_sizes"], [1, 1, self.scheme.group_size]) self.assertEqual(call_kwargs["scale_dtype"], FLOAT8_E8M0FNU_DTYPE) - self.assertEqual(call_kwargs["pertoken_scale_dtype"], FLOAT8_E8M0FNU_DTYPE) self.assertEqual(call_kwargs["output_dtype"], torch.float16) @@ -141,45 +100,28 @@ def setUp(self, mock_ascend, mock_vllm, mock_ep, mock_ensure): mock_ensure.return_value = None self.scheme = AscendW8A8MXFP8DynamicFusedMoEMethod() - def test_quant_type_is_mxfp8(self): - self.assertEqual(self.scheme.quant_type, QuantType.MXFP8) - - def test_model_dtype_is_none(self): - self.assertIsNone(self.scheme.model_dtype) - - 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.float8_e4m3fn) - self.assertEqual(result["w2_weight"].dtype, torch.float8_e4m3fn) - def test_get_weight_various_expert_counts(self): for num_experts in [4, 8, 16]: result = self.scheme.get_weight(num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16) self.assertEqual(result["w13_weight"].shape[0], num_experts) + self.assertEqual(result["w2_weight"].dtype, torch.float8_e4m3fn) def test_get_dynamic_quant_param_dtype_uint8(self): result = self.scheme.get_dynamic_quant_param( self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16 ) - self.assertEqual(result["w13_weight_scale"].dtype, torch.uint8) + self.assertEqual(result["w13_weight_scale"].shape, [8, 512, 4]) self.assertEqual(result["w2_weight_scale"].dtype, torch.uint8) def test_process_weights_stores_original_shapes(self): layer = create_mxfp_moe_layer( num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size ) + original_shape = layer.w13_weight.shape self.scheme.process_weights_after_loading(layer) self.assertTrue(hasattr(layer, "_mxfp8_original_shapes")) self.assertIn("w13_weight", layer._mxfp8_original_shapes) - - def test_process_weights_double_call_idempotent(self): - layer = create_mxfp_moe_layer( - num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size - ) - self.scheme.process_weights_after_loading(layer) - weight_after = layer.w13_weight.data.clone() - self.scheme.process_weights_after_loading(layer) - self.assertTrue(torch.equal(layer.w13_weight.data, weight_after)) + self.assertEqual(layer.w13_weight.shape, (original_shape[0], original_shape[2], original_shape[1])) def test_restore_weights_for_rl_loading(self): layer = create_mxfp_moe_layer( @@ -191,22 +133,6 @@ def test_restore_weights_for_rl_loading(self): self.scheme.restore_weights_for_rl_loading(layer) self.assertEqual(layer.w13_weight.shape, original_w13_shape) - def test_restore_without_transform_is_noop(self): - layer = create_mxfp_moe_layer( - num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size - ) - original_weight = layer.w13_weight.data.clone() - self.scheme.restore_weights_for_rl_loading(layer) - self.assertTrue(torch.equal(layer.w13_weight.data, original_weight)) - - def test_process_weights_transposes_weights(self): - layer = create_mxfp_moe_layer( - num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size - ) - original_shape = layer.w13_weight.shape - self.scheme.process_weights_after_loading(layer) - self.assertEqual(layer.w13_weight.shape, (original_shape[0], original_shape[2], original_shape[1])) - @patch("vllm_ascend.quantization.methods.w8a8_mxfp8._EXTRA_CTX") @patch("vllm_ascend.quantization.methods.w8a8_mxfp8.select_experts") def test_apply_full_params(self, mock_select, mock_ctx): @@ -224,7 +150,7 @@ def test_apply_full_params(self, mock_select, mock_ctx): 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() - result = self.scheme.apply( + self.scheme.apply( layer, x, router_logits, diff --git a/tests/ut/quantization/methods/test_w8a8_pdmix.py b/tests/ut/quantization/methods/test_w8a8_pdmix.py index fe4fe9f763be..4ee87e56d832 100644 --- a/tests/ut/quantization/methods/test_w8a8_pdmix.py +++ b/tests/ut/quantization/methods/test_w8a8_pdmix.py @@ -5,7 +5,6 @@ from tests.ut.base import TestBase from tests.ut.quantization.conftest_quantization import create_mock_vllm_config from vllm_ascend.quantization.methods import ( - AscendW8A8DynamicLinearMethod, AscendW8A8LinearMethod, AscendW8A8PDMixFusedMoeMethod, AscendW8A8PDMixLinearMethod, @@ -16,14 +15,6 @@ class TestAscendW8A8PDMixLinearScheme(TestBase): def setUp(self): self.method = AscendW8A8LinearMethod() - @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") - def test_init_creates_static_and_dynamic_methods(self, mock_vllm_config): - mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None) - scheme = AscendW8A8PDMixLinearMethod() - self.assertTrue(isinstance(scheme._static_method, AscendW8A8LinearMethod)) - self.assertTrue(isinstance(scheme._dynamic_method, AscendW8A8DynamicLinearMethod)) - self.assertFalse(scheme._is_kv_consumer) - @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") @@ -36,7 +27,7 @@ def test_get_weight_delegates_to_static(self, mock_vllm_config, mock_dynamic_cls mock_static_cls.return_value = mock_static_instance scheme = AscendW8A8PDMixLinearMethod() for input_size, output_size in [(64, 128), (256, 512), (1024, 2048)]: - result = scheme.get_weight(input_size, output_size, torch.bfloat16) + scheme.get_weight(input_size, output_size, torch.bfloat16) mock_static_instance.get_weight.assert_called_with(input_size, output_size, torch.bfloat16) mock_dynamic_instance.get_weight.assert_not_called() @@ -51,7 +42,7 @@ def test_get_pertensor_param_delegates_to_static(self, mock_vllm_config, mock_dy mock_static_instance.get_pertensor_param.return_value = {} mock_static_cls.return_value = mock_static_instance scheme = AscendW8A8PDMixLinearMethod() - result = scheme.get_pertensor_param(torch.bfloat16) + scheme.get_pertensor_param(torch.bfloat16) mock_static_instance.get_pertensor_param.assert_called_once_with(torch.bfloat16) mock_dynamic_instance.get_pertensor_param.assert_not_called() @@ -66,7 +57,7 @@ def test_get_perchannel_param_delegates_to_static(self, mock_vllm_config, mock_d mock_static_instance.get_perchannel_param.return_value = {} mock_static_cls.return_value = mock_static_instance scheme = AscendW8A8PDMixLinearMethod() - result = scheme.get_perchannel_param(128, torch.bfloat16) + scheme.get_perchannel_param(128, torch.bfloat16) mock_static_instance.get_perchannel_param.assert_called_once_with(128, torch.bfloat16) mock_dynamic_instance.get_perchannel_param.assert_not_called() @@ -125,21 +116,6 @@ def test_process_weights_after_loading_sets_is_kv_consumer( mock_dynamic_instance.process_weights_after_loading.assert_not_called() self.assertFalse(layer.is_kv_consumer) - @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod") - @patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod") - @patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config") - def test_is_kv_consumer_when_kv_transfer_config_set(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls): - mock_kv_config = MagicMock() - mock_kv_config.is_kv_consumer = True - mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=mock_kv_config) - mock_static_instance = MagicMock() - mock_static_cls.return_value = mock_static_instance - scheme = AscendW8A8PDMixLinearMethod() - layer = MagicMock() - layer.weight_scale = MagicMock(data=torch.randn(128, 1, dtype=torch.bfloat16)) - scheme.process_weights_after_loading(layer) - self.assertTrue(layer.is_kv_consumer) - class TestAscendW8A8PDMixMoEScheme(TestBase): @patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_mc2_group") @@ -159,14 +135,5 @@ def test_get_dynamic_quant_param(self, mock_ascend, mock_vllm, mock_ep, mock_mc2 self.assertEqual(result["w2_deq_scale"].shape, (num_experts, hidden_sizes)) self.assertEqual(result["w2_deq_scale"].dtype, torch.float32) self.assertEqual(result["w13_deq_scale"].shape, (num_experts, 2 * intermediate_size_per_partition)) - self.assertEqual(result["w13_deq_scale"].dtype, torch.float32) - self.assertEqual(result["w2_input_offset"].shape, (num_experts, 1)) self.assertEqual(result["w2_input_offset"].dtype, torch.int8) self.assertEqual(result["w13_input_offset"].shape, (num_experts, 1)) - self.assertEqual(result["w13_input_offset"].dtype, torch.int8) - - # test parent params - self.assertIn("w13_weight_scale", result) - self.assertIn("w13_weight_offset", result) - self.assertIn("w2_weight_scale", result) - self.assertIn("w2_weight_offset", result) diff --git a/tests/ut/quantization/methods/test_w8a8_static.py b/tests/ut/quantization/methods/test_w8a8_static.py index 5c3e9ecaab82..145b43407c10 100644 --- a/tests/ut/quantization/methods/test_w8a8_static.py +++ b/tests/ut/quantization/methods/test_w8a8_static.py @@ -100,30 +100,6 @@ def test_apply_with_x_is_int8(self, mock_npu_quant_matmul, mock_quantize): call_kwargs = mock_npu_quant_matmul.call_args.kwargs self.assertTrue(torch.equal(call_kwargs["bias"], bias)) - @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "0"}) - @patch("torch_npu.npu_format_cast") - def test_process_weights_after_loading_with_nz0(self, mock_npu_format_cast): - layer = MagicMock() - - layer.weight.data = torch.randint(-128, 127, (128, 256), dtype=torch.int8) - layer.input_scale.data = torch.tensor([0.1]) - layer.input_offset.data = torch.tensor([0]) - layer.weight_scale.data = torch.randn(128, 1) - layer.weight_offset.data = torch.randn(128, 1) - - mock_npu_format_cast.side_effect = identity - self.method.process_weights_after_loading(layer) - - expected_offset = torch.tensor([0]).repeat(256).to(torch.int8) - self.assertTrue(torch.equal(layer.aclnn_input_offset.data, expected_offset)) - self.assertFalse(layer.aclnn_input_offset.requires_grad) - - self.assertEqual(layer.weight.data.shape, (256, 128)) - self.assertEqual(layer.weight_scale.data.shape, (128,)) - self.assertEqual(layer.weight_offset.data.shape, (128,)) - mock_npu_format_cast.assert_not_called() - self.assertNotIn("deq_scale", dict(layer.named_parameters())) - @patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_NZ": "1"}) @patch("torch_npu.npu_format_cast") def test_process_weights_after_loading_with_nz1(self, mock_npu_format_cast): @@ -171,5 +147,4 @@ def test_process_weights_after_loading_with_nz2_and_compressed_tensors(self, moc self.assertEqual(layer.weight_scale.data.shape, (128,)) self.assertEqual(layer.weight_offset.data.shape, (128,)) mock_npu_format_cast.assert_called_once() - self.assertIn("deq_scale", dict(layer.named_parameters())) self.assertFalse(isinstance(layer.deq_scale, MagicMock)) diff --git a/tests/ut/quantization/test_compressed_tensors_config.py b/tests/ut/quantization/test_compressed_tensors_config.py index b24a660924db..617329319ae2 100644 --- a/tests/ut/quantization/test_compressed_tensors_config.py +++ b/tests/ut/quantization/test_compressed_tensors_config.py @@ -1,6 +1,5 @@ from unittest.mock import MagicMock, patch -import torch from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.linear import RowParallelLinear, UnquantizedLinearMethod @@ -14,46 +13,6 @@ from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD -class TestAscendCompressedTensorsConfigBasic(TestBase): - def test_get_name(self): - config = AscendCompressedTensorsConfig.from_config(COMPRESSED_TENSORS_W8A8_CONFIG) - self.assertEqual(config.get_name(), "compressed-tensors") - - def test_get_supported_act_dtypes(self): - dtypes = AscendCompressedTensorsConfig.get_supported_act_dtypes() - self.assertIn(torch.int8, dtypes) - self.assertIn(torch.float16, dtypes) - self.assertIn(torch.bfloat16, dtypes) - self.assertEqual(len(dtypes), 3) - - def test_get_min_capability_raises(self): - with self.assertRaises(NotImplementedError): - AscendCompressedTensorsConfig.get_min_capability() - - def test_get_config_filenames(self): - filenames = AscendCompressedTensorsConfig.get_config_filenames() - self.assertEqual(filenames, []) - - def test_init(self): - config = AscendCompressedTensorsConfig.from_config(COMPRESSED_TENSORS_W8A8_CONFIG) - self.assertEqual(config.ignore, ["lm_head"]) - self.assertEqual(config.quant_format, "int-quantized") - self.assertEqual(list(config.target_scheme_map.keys()), ["Linear"]) - self.assertEqual(config.target_scheme_map["Linear"]["format"], "int-quantized") - self.assertIsNotNone(config.target_scheme_map["Linear"]["input_activations"]) - - def test_apply_vllm_mapper(self): - hf_to_vllm_mapper = MagicMock() - config = AscendCompressedTensorsConfig( - target_scheme_map={"Linear": {}}, - ignore=["lm_head"], - quant_format="", - ) - config.apply_vllm_mapper(hf_to_vllm_mapper) - hf_to_vllm_mapper.apply_dict.assert_called_once() - hf_to_vllm_mapper.apply_list.assert_called_once() - - class TestAscendCompressedTensorsQuanType(TestBase): def setUp(self): self.config = AscendCompressedTensorsConfig( @@ -120,7 +79,9 @@ class TestAscendCompressedTensorsConfigGetQuantMethod(TestBase): def setUp(self): self.config = AscendCompressedTensorsConfig.from_config(COMPRESSED_TENSORS_W8A8_CONFIG) - def test_get_linear_quant_method(self): + @patch("vllm_ascend.quantization.method_adapters.AscendLinearMethod.__init__") + def test_get_linear_quant_method(self, mock_method): + mock_method.return_value = None layer = MagicMock(spec=RowParallelLinear) result = self.config.get_quant_method(layer, "model.layers.0.self_attn.q_proj") self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) @@ -133,12 +94,9 @@ def test_get_linear_unquantized_method(self): self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) self.assertTrue(isinstance(result, UnquantizedLinearMethod)) - from vllm_ascend.quantization.methods import AscendW8A8DynamicFusedMoEMethod, AscendW8A8DynamicLinearMethod - @patch("vllm_ascend.quantization.methods.AscendW8A8DynamicFusedMoEMethod.__init__") def test_get_moe_quant_method(self, mock_method): mock_method.return_value = None - mock_method.return_value = MagicMock(spec=AscendW8A8DynamicFusedMoEMethod) layer = MagicMock(spec=FusedMoE) layer.moe_config = {} result = self.config.get_quant_method(layer, "model.layers.0.mlp.experts") diff --git a/tests/ut/quantization/test_method_adapters.py b/tests/ut/quantization/test_method_adapters.py index d7604df18423..263fb9ca08aa 100644 --- a/tests/ut/quantization/test_method_adapters.py +++ b/tests/ut/quantization/test_method_adapters.py @@ -6,7 +6,6 @@ from tests.ut.base import TestBase from vllm_ascend.quantization.method_adapters import ( - AscendEmbeddingMethod, AscendFusedMoEMethod, AscendKVCacheMethod, AscendLinearMethod, @@ -15,6 +14,7 @@ class TestAscendLinearMethod(TestBase): + @patch("vllm_ascend.quantization.method_adapters.enable_dsa_cp_with_layer_shard") def setUp(self): self.mock_scheme = MagicMock(spec=AscendLinearScheme) self.mock_scheme.get_weight.return_value = { @@ -82,19 +82,6 @@ def test_process_weights_after_loading_delegates(self): self.method.process_weights_after_loading(layer) self.mock_scheme.process_weights_after_loading.assert_called_once_with(layer) - def test_process_weights_after_loading_skips_when_not_present(self): - del self.mock_scheme.process_weights_after_loading - layer = torch.nn.Module() - self.method.process_weights_after_loading(layer) - - def test_get_computed_params(self): - computed = self.method.get_computed_params() - self.assertIn("weight_offset", computed) - self.assertIn("quant_bias", computed) - self.assertIn("deq_scale", computed) - self.assertIn("weight_scale", computed) - self.assertEqual(len(computed), 4) - def test_apply_delegates_to_scheme(self): layer = MagicMock(spec=ColumnParallelLinear) x = torch.randn(4, 256) @@ -154,11 +141,6 @@ def test_process_weights_after_loading_delegates(self): self.method.process_weights_after_loading(layer) self.mock_scheme.process_weights_after_loading.assert_called_once_with(layer) - def test_process_weights_after_loading_skips_when_not_present(self): - del self.mock_scheme.process_weights_after_loading - layer = torch.nn.Module() - self.method.process_weights_after_loading(layer) - def test_create_weights_registers_parameters(self): self.mock_scheme.get_weight.return_value = { "w13_weight": torch.empty(8, 256, 128, dtype=torch.int8), @@ -172,14 +154,7 @@ def test_create_weights_registers_parameters(self): "w2_weight_offset": torch.empty(8, 128, 1, dtype=torch.bfloat16), } # per channel quantization - layer = torch.nn.Module() - self.method.create_weights( - layer, - num_experts=8, - hidden_size=128, - intermediate_size_per_partition=256, - params_dtype=torch.bfloat16, - ) + layer = self.create_moe_weights() self.assertIn("w13_weight", dict(layer.named_parameters())) self.assertIn("w2_weight", dict(layer.named_parameters())) @@ -191,6 +166,11 @@ def test_create_weights_registers_parameters(self): # per group quantization self.mock_scheme.group_size = 128 + layer = self.create_moe_weights() + self.assertEqual(layer.w13_weight_scale.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + self.assertEqual(layer.w2_weight_offset.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + + def create_moe_weights(self): layer = torch.nn.Module() self.method.create_weights( layer, @@ -199,11 +179,7 @@ def test_create_weights_registers_parameters(self): intermediate_size_per_partition=256, params_dtype=torch.bfloat16, ) - self.assertEqual(layer.w13_weight_scale_second.quant_method, FusedMoeWeightScaleSupported.GROUP.value) - self.assertEqual(layer.w2_weight_offset_second.quant_method, FusedMoeWeightScaleSupported.GROUP.value) - self.assertEqual(layer.w2_scale_bias.quant_method, FusedMoeWeightScaleSupported.GROUP.value) - self.assertEqual(layer.w13_weight_scale.quant_method, FusedMoeWeightScaleSupported.GROUP.value) - self.assertEqual(layer.w2_weight_offset.quant_method, FusedMoeWeightScaleSupported.GROUP.value) + return layer def test_apply_method(self): layer = torch.nn.Module() @@ -214,18 +190,3 @@ def test_apply_method(self): self.mock_scheme.apply.return_value = None self.method.apply(layer, x, router_logits, top_k, renormalize) self.mock_scheme.apply.assert_called_once() - - def test_supports_eplb_default_false(self): - self.assertFalse(self.method.supports_eplb) - - def test_supports_eplb_when_scheme_supports(self): - self.mock_scheme.supports_eplb = True - self.assertTrue(self.method.supports_eplb) - - -class TestAscendEmbeddingMethod(TestBase): - def test_init(self): - layer = MagicMock(spec=AscendLinearScheme) - method = AscendEmbeddingMethod(layer) - self.assertTrue(isinstance(method, AscendLinearMethod)) - self.assertEqual(method.quant_method, layer) diff --git a/tests/ut/quantization/test_modelslim_config.py b/tests/ut/quantization/test_modelslim_config.py index a5d563a57ec9..c8ddbb4087a0 100644 --- a/tests/ut/quantization/test_modelslim_config.py +++ b/tests/ut/quantization/test_modelslim_config.py @@ -15,10 +15,6 @@ from vllm_ascend.quantization.modelslim_config import ( MODELSLIM_CONFIG_FILENAME, AscendModelSlimConfig, - create_scheme_for_layer, - get_linear_quant_type, - get_packed_modules_mapping, - get_quant_type_for_layer, ) from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD @@ -203,10 +199,6 @@ def test_is_layer_skipped_ascend(self): with self.assertRaises(ValueError): config.is_layer_skipped_ascend("fused_layer", fused_mapping) - def test_init_with_none_config(self): - config = AscendModelSlimConfig(None) - self.assertEqual(config.quant_description, {}) - def test_init_with_default_config(self): config = AscendModelSlimConfig() self.assertEqual(config.quant_description, {}) @@ -286,121 +278,6 @@ def test_apply_extra_quant_adaptations_weight_packed(self): self.assertEqual(config.quant_description["model.layers.0.weight"], "INT8") -class TestGetPackedModulesMapping(TestBase): - def test_known_model_type(self): - result = get_packed_modules_mapping("deepseek_v3") - self.assertIsInstance(result, dict) - self.assertIn("gate_up_proj", result) - self.assertIn("experts", result) - - def test_unknown_model_type(self): - result = get_packed_modules_mapping("unknown_model") - self.assertEqual(result, {}) - - -class TestGetLinearQuantType(TestBase): - def test_simple_non_fused_layer(self): - quant_description = {"layer1.weight": "W8A8_DYNAMIC"} - packed_modules_mapping = {} - result = get_linear_quant_type(quant_description, "layer1", packed_modules_mapping) - self.assertEqual(result, "W8A8_DYNAMIC") - - def test_fused_layer_consistent_quant_types(self): - quant_description = { - "shard1.weight": "W8A8_DYNAMIC", - "shard2.weight": "W8A8_DYNAMIC", - } - packed_modules_mapping = {"fused_layer": ["shard1", "shard2"]} - result = get_linear_quant_type(quant_description, "fused_layer", packed_modules_mapping) - self.assertEqual(result, "W8A8_DYNAMIC") - - def test_fused_layer_inconsistent_quant_types(self): - quant_description = { - "shard1.weight": "W8A8_DYNAMIC", - "shard2.weight": "FLOAT", - } - packed_modules_mapping = {"fused_layer": ["shard1", "shard2"]} - with self.assertRaises(ValueError) as ctx: - get_linear_quant_type(quant_description, "fused_layer", packed_modules_mapping) - self.assertIn("Not all shards", str(ctx.exception)) - - -class TestGetQuantTypeForLayer(TestBase): - def test_attention_layer_with_fa_quant_type(self): - quant_description = { - "fa_quant_type": "C8", - "layer1.weight": "W8A8_DYNAMIC", - } - result = get_quant_type_for_layer(quant_description, "layer1", "attention") - self.assertEqual(result, "C8") - - def test_attention_layer_with_indexer_quant_type(self): - quant_description = { - "indexer_quant_type": "INT8", - "layer1.weight": "W8A8_DYNAMIC", - } - result = get_quant_type_for_layer(quant_description, "layer1", "attention") - self.assertEqual(result, "INT8") - - def test_attention_layer_with_layer_indexer_quant_type(self): - quant_description = { - "layer1.indexer.quant_type": "INT8", - "layer1.weight": "W8A8_DYNAMIC", - } - result = get_quant_type_for_layer(quant_description, "layer1", "attention") - self.assertEqual(result, "INT8") - - def test_linear_layer(self): - quant_description = { - "layer1.weight": "W8A8_DYNAMIC", - } - result = get_quant_type_for_layer(quant_description, "layer1", "linear") - self.assertEqual(result, "W8A8_DYNAMIC") - - def test_moe_layer(self): - quant_description = { - "moe_layer.weight": "W8A8_DYNAMIC", - } - result = get_quant_type_for_layer(quant_description, "moe_layer", "moe") - self.assertEqual(result, "W8A8_DYNAMIC") - - def test_with_none_packed_modules_mapping(self): - quant_description = { - "layer1.weight": "W8A8_DYNAMIC", - } - result = get_quant_type_for_layer(quant_description, "layer1", "linear", None) - self.assertEqual(result, "W8A8_DYNAMIC") - - -class TestCreateSchemeForLayer(TestBase): - @patch("vllm_ascend.quantization.modelslim_config.get_scheme_class") - def test_successful_scheme_creation(self, mock_get_scheme_class): - mock_scheme_class = MagicMock() - mock_scheme_instance = MagicMock() - mock_scheme_class.return_value = mock_scheme_instance - mock_get_scheme_class.return_value = mock_scheme_class - - quant_description = {"layer1.weight": "W8A8_DYNAMIC"} - result = create_scheme_for_layer(quant_description, "layer1", "linear") - - self.assertIs(result, mock_scheme_instance) - mock_get_scheme_class.assert_called_once_with("W8A8_DYNAMIC", "linear") - - def test_quant_type_none_raises_error(self): - quant_description = {} - with self.assertRaises(KeyError) as ctx: - create_scheme_for_layer(quant_description, "layer1", "linear") - - @patch("vllm_ascend.quantization.modelslim_config.get_scheme_class") - def test_unsupported_quant_type_raises_error(self, mock_get_scheme_class): - mock_get_scheme_class.return_value = None - - quant_description = {"layer1.weight": "UNSUPPORTED_TYPE"} - with self.assertRaises(NotImplementedError) as ctx: - create_scheme_for_layer(quant_description, "layer1", "linear") - self.assertIn("doesn't support UNSUPPORTED_TYPE", str(ctx.exception)) - - class TestApplyVllmMapper(TestBase): def test_apply_mapper_with_populated_quant_description(self): config = AscendModelSlimConfig({"old_key.weight": "INT8"}) @@ -412,189 +289,27 @@ def test_apply_mapper_with_populated_quant_description(self): self.assertEqual(config.quant_description, {"new_key.weight": "INT8"}) mock_mapper.apply_dict.assert_called_once_with({"old_key.weight": "INT8"}) - def test_idempotency_when_mapper_already_applied(self): - config = AscendModelSlimConfig({"key.weight": "INT8"}) - mock_mapper = MagicMock() - mock_mapper.apply_dict.return_value = {"key.weight": "INT8"} - - config.apply_vllm_mapper(mock_mapper) - first_result = config.quant_description - - config.apply_vllm_mapper(mock_mapper) - second_result = config.quant_description - - self.assertEqual(first_result, second_result) - mock_mapper.apply_dict.assert_called_once() - - def test_empty_quant_description(self): - config = AscendModelSlimConfig({}) - mock_mapper = MagicMock() - - config.apply_vllm_mapper(mock_mapper) - - mock_mapper.apply_dict.assert_not_called() - class TestGetCacheScale(TestBase): def test_c8_kv_cache_type_k_proj_scale(self): config = AscendModelSlimConfig({"kv_cache_type": "C8"}) result = config.get_cache_scale("model.layers.0.k_proj.kv_cache_scale") self.assertEqual(result, "model.layers.0.attn.k_cache_scale") - - def test_c8_kv_cache_type_k_proj_offset(self): - config = AscendModelSlimConfig({"kv_cache_type": "C8"}) - result = config.get_cache_scale("model.layers.0.k_proj.kv_cache_offset") - self.assertEqual(result, "model.layers.0.attn.k_cache_offset") - - def test_c8_kv_cache_type_v_proj_scale(self): - config = AscendModelSlimConfig({"kv_cache_type": "C8"}) - result = config.get_cache_scale("model.layers.0.v_proj.kv_cache_scale") - self.assertEqual(result, "model.layers.0.attn.v_cache_scale") - - def test_c8_kv_cache_type_v_proj_offset(self): - config = AscendModelSlimConfig({"kv_cache_type": "C8"}) result = config.get_cache_scale("model.layers.0.v_proj.kv_cache_offset") self.assertEqual(result, "model.layers.0.attn.v_cache_offset") - def test_non_c8_kv_cache_type(self): + def test_no_match(self): config = AscendModelSlimConfig({"kv_cache_type": "FLOAT"}) result = config.get_cache_scale("model.layers.0.k_proj.kv_cache_scale") self.assertIsNone(result) - def test_no_matching_suffix(self): config = AscendModelSlimConfig({"kv_cache_type": "C8"}) result = config.get_cache_scale("model.layers.0.other_key") self.assertIsNone(result) -class TestQuantPrefixMapper(TestBase): - def test_sets_model_type_and_returns_prefix(self): - config = AscendModelSlimConfig({}) - result = config.quant_prefix_mapper("deepseek_v3", "model.layers.0") - self.assertEqual(result, "model.layers.0") - self.assertEqual(config.model_type, "deepseek_v3") - - -class TestIsFaQuantLayer(TestBase): - def test_enable_fa_quant_true_matching_layer(self): - config = AscendModelSlimConfig( - { - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - } - ) - result = config.is_fa_quant_layer("layers.1.attn") - self.assertTrue(result) - - def test_enable_fa_quant_false(self): - config = AscendModelSlimConfig({}) - result = config.is_fa_quant_layer("layers.1.attn") - self.assertFalse(result) - - def test_non_matching_layer_index(self): - config = AscendModelSlimConfig( - { - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - } - ) - result = config.is_fa_quant_layer("layers.2.attn") - self.assertFalse(result) - - def test_no_layer_index(self): - config = AscendModelSlimConfig( - { - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - } - ) - result = config.is_fa_quant_layer("model.attn") - self.assertFalse(result) - - -class TestEnablingFaQuant(TestBase): - def test_decode_instance_with_fa_quant_layer(self): - config = AscendModelSlimConfig( - { - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - } - ) - mock_vllm_config = MagicMock() - mock_vllm_config.kv_transfer_config.is_kv_consumer = True - mock_vllm_config.kv_transfer_config.is_kv_producer = False - - result = config.enabling_fa_quant(mock_vllm_config, "layers.1.attn") - self.assertTrue(result) - - def test_non_decode_instance(self): - config = AscendModelSlimConfig( - { - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - } - ) - mock_vllm_config = MagicMock() - mock_vllm_config.kv_transfer_config.is_kv_consumer = False - mock_vllm_config.kv_transfer_config.is_kv_producer = True - - result = config.enabling_fa_quant(mock_vllm_config, "layers.1.attn") - self.assertFalse(result) - - def test_is_fa_quant_layer_false(self): - config = AscendModelSlimConfig( - { - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - } - ) - mock_vllm_config = MagicMock() - mock_vllm_config.kv_transfer_config.is_kv_consumer = True - mock_vllm_config.kv_transfer_config.is_kv_producer = False - - result = config.enabling_fa_quant(mock_vllm_config, "layers.2.attn") - self.assertFalse(result) - - -class TestIsIndexerQuantLayer(TestBase): - def test_enable_indexer_quant_true_matching_layer(self): - config = AscendModelSlimConfig( - { - "indexer_quant_type": "INT8", - "layers.1.indexer.quant_type": "INT8", - } - ) - result = config.is_indexer_quant_layer("layers.1.attn") - self.assertTrue(result) - - def test_enable_indexer_quant_false(self): - config = AscendModelSlimConfig({}) - result = config.is_indexer_quant_layer("layers.1.attn") - self.assertFalse(result) - - def test_non_matching_layer_index(self): - config = AscendModelSlimConfig( - { - "indexer_quant_type": "INT8", - "layers.1.indexer.quant_type": "INT8", - } - ) - result = config.is_indexer_quant_layer("layers.2.attn") - self.assertFalse(result) - - def test_no_layer_index(self): - config = AscendModelSlimConfig( - { - "indexer_quant_type": "INT8", - "layers.1.indexer.quant_type": "INT8", - } - ) - result = config.is_indexer_quant_layer("model.attn") - self.assertFalse(result) - - class TestGetKvQuantDtype(TestBase): - def test_enable_fa_quant_true_with_mla(self): + def test_enable_fa_quant(self): config = AscendModelSlimConfig( { "fa_quant_type": "C8", @@ -603,23 +318,14 @@ def test_enable_fa_quant_true_with_mla(self): ) mock_model_config = MagicMock() mock_model_config.dtype = torch.float16 + # test mla mock_model_config.use_mla = True - k_dtype, v_dtype = config.get_kv_quant_dtype("layers.1.attn", torch.float16, mock_model_config) self.assertEqual(k_dtype, torch.int8) self.assertEqual(v_dtype, torch.float16) - def test_enable_fa_quant_true_without_mla(self): - config = AscendModelSlimConfig( - { - "fa_quant_type": "C8", - "layers.1.fa_k.scale": "C8", - } - ) - mock_model_config = MagicMock() - mock_model_config.dtype = torch.float16 + # test gqa mock_model_config.use_mla = False - k_dtype, v_dtype = config.get_kv_quant_dtype("layers.1.attn", torch.float16, mock_model_config) self.assertEqual(k_dtype, torch.int8) self.assertEqual(v_dtype, torch.int8) @@ -628,11 +334,8 @@ def test_enable_fa_quant_false(self): config = AscendModelSlimConfig({}) mock_model_config = MagicMock() mock_model_config.dtype = torch.float16 - mock_model_config.use_mla = True - k_dtype, v_dtype = config.get_kv_quant_dtype("layers.1.attn", torch.float16, mock_model_config) self.assertEqual(k_dtype, torch.float16) - self.assertEqual(v_dtype, torch.float16) class TestGetKvQuantSplitFactor(TestBase): @@ -675,7 +378,7 @@ def test_with_fa_quant_type(self): self.assertTrue(config.enable_fa_quant) self.assertIn(1, config.kvcache_quant_layers) - self.assertIn(2, config.kvcache_quant_layers) + self.assertNotIn(5, config.kvcache_quant_layers) self.assertFalse(config.enable_indexer_quant) self.assertEqual(config.indexer_quant_layers, []) @@ -693,7 +396,7 @@ def test_with_indexer_quant_type(self): self.assertEqual(config.kvcache_quant_layers, []) self.assertTrue(config.enable_indexer_quant) self.assertIn(1, config.indexer_quant_layers) - self.assertIn(3, config.indexer_quant_layers) + self.assertNotIn(5, config.indexer_quant_layers) def test_with_neither_quant_type(self): config = AscendModelSlimConfig({}) diff --git a/tests/ut/quantization/test_quant_parser.py b/tests/ut/quantization/test_quant_parser.py index e6ca5de1f583..c061aaded455 100644 --- a/tests/ut/quantization/test_quant_parser.py +++ b/tests/ut/quantization/test_quant_parser.py @@ -1,40 +1,13 @@ from unittest.mock import patch -import torch - from tests.ut.base import TestBase from vllm_ascend.quantization.quant_parser import ( - QuantTypeMapping, get_rollback_quant_type, parse_mxfp_quant_params, parse_quant_moe_down_proj_params, ) -class TestQuantTypeMapping(TestBase): - def test_get_quant_settings_returns_dict(self): - settings = QuantTypeMapping.get_quant_settings() - self.assertIsInstance(settings, dict) - - def test_get_quant_settings_contains_expected_keys(self): - settings = QuantTypeMapping.get_quant_settings() - self.assertIn("W8A8_MXFP8", settings) - self.assertIn("W4A4_MXFP4", settings) - self.assertIn("W4A8_MXFP", settings) - - def test_w8a8_mxfp8_settings(self): - settings = QuantTypeMapping.get_quant_settings()["W8A8_MXFP8"] - self.assertEqual(settings["act_quant_type"], torch.float8_e4m3fn) - self.assertIsNone(settings["weight_quant_type"]) - - def test_w4a4_mxfp4_settings(self): - settings = QuantTypeMapping.get_quant_settings()["W4A4_MXFP4"] - self.assertIn("act_quant_type", settings) - self.assertIn("weight_quant_type", settings) - self.assertIn("scale_dtype", settings) - self.assertIn("per_token_scale_dtype", settings) - - class TestGetRollbackQuantType(TestBase): def test_returns_down_proj_quant_type(self): config = { @@ -49,42 +22,13 @@ def test_returns_default_when_no_down_proj(self): result = get_rollback_quant_type(config) self.assertEqual(result, "W8A8_MXFP8") - def test_returns_down_proj_type_with_multiple_entries(self): - config = { - "model.layers.0.mlp.gate_proj": "W8A8_MXFP8", - "model.layers.0.mlp.up_proj": "W8A8_MXFP8", - "model.layers.0.mlp.down_proj": "W4A8_MXFP", - } - result = get_rollback_quant_type(config) - self.assertEqual(result, "W4A8_MXFP") - class TestParseMxfpQuantParams(TestBase): def test_default_values(self): act, weight, scale, per_token, round_mode = parse_mxfp_quant_params() - self.assertEqual(act, torch.float8_e4m3fn) - self.assertEqual(weight, torch.float8_e4m3fn) - self.assertIsNone(scale) - self.assertIsNone(per_token) - self.assertEqual(round_mode, "rint") - - def test_custom_values(self): - act, weight, scale, per_token, round_mode = parse_mxfp_quant_params( - act_quant_type=torch.float16, - weight_quant_type=torch.float8_e4m3fn, - round_mode="round", - ) - self.assertEqual(act, torch.float16) - self.assertEqual(weight, torch.float8_e4m3fn) - self.assertEqual(round_mode, "round") - - def test_scale_type_overrides(self): - act, weight, scale, per_token, round_mode = parse_mxfp_quant_params( - scale_type="float8_e8m0fnu", - per_token_scale_type="float8_e8m0fnu", - ) - self.assertEqual(scale, "float8_e8m0fnu") - self.assertEqual(per_token, "float8_e8m0fnu") + self.assertIsNotNone(act) + self.assertIsNotNone(weight) + self.assertIsNotNone(round_mode) class TestParseQuantMoeDownProjParams(TestBase): @@ -99,15 +43,3 @@ def test_w4a4_mxfp4_respects_parsed_round_mode(self, mock_ensure): mock_ensure.return_value = None act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params("W4A4_MXFP4", "round") self.assertEqual(round_mode, "round") - - @patch("vllm_ascend.quantization.quant_parser.ensure_mxfp4_dtype_available") - def test_w4a4_mxfp4_rint_round_mode(self, mock_ensure): - mock_ensure.return_value = None - act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params("W4A4_MXFP4", "rint") - self.assertEqual(round_mode, "rint") - - @patch("vllm_ascend.quantization.quant_parser.ensure_mxfp8_scale_dtype_available") - def test_w4a8_mxfp_uses_rint_round_mode(self, mock_ensure): - mock_ensure.return_value = None - act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params("W4A8_MXFP", "round") - self.assertEqual(round_mode, "rint") diff --git a/tests/ut/quantization/test_quant_type.py b/tests/ut/quantization/test_quant_type.py deleted file mode 100644 index c70ab7298815..000000000000 --- a/tests/ut/quantization/test_quant_type.py +++ /dev/null @@ -1,28 +0,0 @@ -from tests.ut.base import TestBase -from vllm_ascend.quantization.quant_type import QuantType - - -class TestQuantType(TestBase): - def test_enum_values(self): - self.assertEqual(QuantType.NONE.value, 0) - self.assertEqual(QuantType.W8A8.value, 1) - self.assertEqual(QuantType.W4A8.value, 2) - self.assertEqual(QuantType.MXFP8.value, 3) - self.assertEqual(QuantType.W4A16.value, 4) - self.assertEqual(QuantType.MXFP4.value, 5) - - def test_enum_members_count(self): - members = list(QuantType) - self.assertEqual(len(members), 6) - - def test_enum_identity(self): - self.assertIs(QuantType.W8A8, QuantType.W8A8) - self.assertIsNot(QuantType.W8A8, QuantType.W4A8) - - def test_enum_by_value(self): - self.assertEqual(QuantType(1), QuantType.W8A8) - self.assertEqual(QuantType(3), QuantType.MXFP8) - - def test_enum_by_name(self): - self.assertEqual(QuantType["NONE"], QuantType.NONE) - self.assertEqual(QuantType["W4A16"], QuantType.W4A16) diff --git a/tests/ut/quantization/test_utils.py b/tests/ut/quantization/test_utils.py index 0682e72c187b..d8dd5948d19e 100644 --- a/tests/ut/quantization/test_utils.py +++ b/tests/ut/quantization/test_utils.py @@ -162,19 +162,18 @@ def test_passes_revision_to_detect(self, mock_detect, mock_get_quant): class TestEnableFaQuant(TestBase): def test_non_quantization_scenarios(self): + # non quantization scene vllm_config = MagicMock() vllm_config.quant_config = None result = enable_fa_quant(vllm_config) self.assertFalse(result) - def test_llmcompressor_quantization_scenario(self): - vllm_config = MagicMock() + # CompressedTensors scene vllm_config.quant_config = AscendCompressedTensorsConfig({}, [], "", {}) result = enable_fa_quant(vllm_config) self.assertFalse(result) - def test_not_fa3_quantization_scenario(self): - vllm_config = MagicMock() + # non fa3 quant scene vllm_config.quant_config = AscendModelSlimConfig(W8A8_CONFIG) result = enable_fa_quant(vllm_config) self.assertFalse(result) diff --git a/vllm_ascend/quantization/method_adapters.py b/vllm_ascend/quantization/method_adapters.py index 7d52f8ca39d2..cd68ddcdd104 100644 --- a/vllm_ascend/quantization/method_adapters.py +++ b/vllm_ascend/quantization/method_adapters.py @@ -112,7 +112,6 @@ def create_weights( or is_mx_quant_type(self.quant_method) ): param.input_dim = 1 - param.input_dim = 1 def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if hasattr(self.quant_method, "process_weights_after_loading"): diff --git a/vllm_ascend/quantization/methods/kv_c8.py b/vllm_ascend/quantization/methods/kv_c8.py index cb351b1074af..34a70fac2f37 100644 --- a/vllm_ascend/quantization/methods/kv_c8.py +++ b/vllm_ascend/quantization/methods/kv_c8.py @@ -25,8 +25,6 @@ def _fa_quant_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor): @register_scheme("FAKQuant", "attention") class AscendFAQuantAttentionMethod: def __init__(self): - self.transpose_weight = True - self.printFlag = False vllm_config = get_current_vllm_config() config = vllm_config.model_config.hf_config self.kv_lora_rank = getattr(config, "kv_lora_rank", 0) diff --git a/vllm_ascend/quantization/methods/w4a16.py b/vllm_ascend/quantization/methods/w4a16.py index 3a0e3f5e3142..38c1ca2a40b9 100644 --- a/vllm_ascend/quantization/methods/w4a16.py +++ b/vllm_ascend/quantization/methods/w4a16.py @@ -107,7 +107,6 @@ class AscendW4A16FusedMoEMethod(AscendMoEScheme): quant_type: QuantType = QuantType.W4A16 def __init__(self) -> None: - self.transpose_weight = True self.num_bits = 4 # dtype = torch.int4 self.pack_factor = 8 # pack 8 of torch.int4 tensors to torch.int32 @@ -246,36 +245,35 @@ def apply( ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - if self.transpose_weight: - w13_shape = layer.w13_weight_packed.data.shape - w2_shape = layer.w2_weight_packed.data.shape - unpacked_w13_weight = ( - unpack_from_int32( - layer.w13_weight_packed.data.flatten(0, 1), - torch.Size([w13_shape[0] * w13_shape[1], w13_shape[2] * self.pack_factor]), - self.num_bits, - ) - .view(w13_shape[0], w13_shape[1], -1) - .transpose(1, 2) - .contiguous() - .int() + w13_shape = layer.w13_weight_packed.data.shape + w2_shape = layer.w2_weight_packed.data.shape + unpacked_w13_weight = ( + unpack_from_int32( + layer.w13_weight_packed.data.flatten(0, 1), + torch.Size([w13_shape[0] * w13_shape[1], w13_shape[2] * self.pack_factor]), + self.num_bits, ) - unpacked_w2_weight = ( - unpack_from_int32( - layer.w2_weight_packed.data.flatten(0, 1), - torch.Size([w2_shape[0] * w2_shape[1], w2_shape[2] * self.pack_factor]), - self.num_bits, - ) - .view(w2_shape[0], w2_shape[1], -1) - .transpose(1, 2) - .contiguous() - .int() + .view(w13_shape[0], w13_shape[1], -1) + .transpose(1, 2) + .contiguous() + .int() + ) + unpacked_w2_weight = ( + unpack_from_int32( + layer.w2_weight_packed.data.flatten(0, 1), + torch.Size([w2_shape[0] * w2_shape[1], w2_shape[2] * self.pack_factor]), + self.num_bits, ) - layer.w13_weight_packed.data = pack_to_int32(unpacked_w13_weight) - layer.w2_weight_packed.data = pack_to_int32(unpacked_w2_weight) + .view(w2_shape[0], w2_shape[1], -1) + .transpose(1, 2) + .contiguous() + .int() + ) + layer.w13_weight_packed.data = pack_to_int32(unpacked_w13_weight) + layer.w2_weight_packed.data = pack_to_int32(unpacked_w2_weight) - layer.w13_weight_scale.data = layer.w13_weight_scale.data.transpose(1, 2).contiguous() - layer.w2_weight_scale.data = layer.w2_weight_scale.data.transpose(1, 2).contiguous() + layer.w13_weight_scale.data = layer.w13_weight_scale.data.transpose(1, 2).contiguous() + layer.w2_weight_scale.data = layer.w2_weight_scale.data.transpose(1, 2).contiguous() - layer.w13_weight_offset.data = layer.w13_weight_offset.data.transpose(1, 2).contiguous() - layer.w2_weight_offset.data = layer.w2_weight_offset.data.transpose(1, 2).contiguous() + layer.w13_weight_offset.data = layer.w13_weight_offset.data.transpose(1, 2).contiguous() + layer.w2_weight_offset.data = layer.w2_weight_offset.data.transpose(1, 2).contiguous() diff --git a/vllm_ascend/quantization/methods/w4a4_laos_dynamic.py b/vllm_ascend/quantization/methods/w4a4_laos_dynamic.py index 455b57fab75f..ded946545fa0 100644 --- a/vllm_ascend/quantization/methods/w4a4_laos_dynamic.py +++ b/vllm_ascend/quantization/methods/w4a4_laos_dynamic.py @@ -34,7 +34,7 @@ class AscendW4A4LaosDynamicLinearMethod(AscendLinearScheme): """ def __init__(self): - self.transpose_weight = True + pass def get_weight(self, input_size: int, output_size: int, params_dtype: torch.dtype) -> dict[str, Any]: params_dict = {"weight": torch.empty(output_size, input_size, dtype=torch.int8)} @@ -72,5 +72,4 @@ def apply( def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight_scale.data = layer.weight_scale.data.to(torch.float32) layer.weight.data = torch_npu.npu_convert_weight_to_int4pack(layer.weight.data.to(torch.int32)) - if self.transpose_weight: - layer.weight.data = layer.weight.data.transpose(-1, -2) + layer.weight.data = layer.weight.data.transpose(-1, -2) From a9a288a3984983efd7a88ee26b969582dc2517ab Mon Sep 17 00:00:00 2001 From: Wang Kunpeng <1289706727@qq.com> Date: Fri, 24 Apr 2026 22:22:34 +0800 Subject: [PATCH 4/4] fix ut Signed-off-by: Wang Kunpeng <1289706727@qq.com> --- tests/ut/quantization/methods/test_w8a8_dynamic.py | 1 - tests/ut/quantization/methods/test_w8a8_mxfp8.py | 2 +- tests/ut/quantization/test_method_adapters.py | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/ut/quantization/methods/test_w8a8_dynamic.py b/tests/ut/quantization/methods/test_w8a8_dynamic.py index e9998b2eaabf..74c00480038d 100644 --- a/tests/ut/quantization/methods/test_w8a8_dynamic.py +++ b/tests/ut/quantization/methods/test_w8a8_dynamic.py @@ -174,6 +174,5 @@ def test_process_weights_after_loading(self, mock_envs, mock_format_cast): num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size ) self.quant_method.process_weights_after_loading(layer) - self.assertEqual(layer.fused_w1_scale.dtype, torch.int64) self.assertTrue(hasattr(layer, "w13_weight_list")) self.assertFalse(hasattr(layer, "w13_weight_scale_fp32")) diff --git a/tests/ut/quantization/methods/test_w8a8_mxfp8.py b/tests/ut/quantization/methods/test_w8a8_mxfp8.py index f46d17bda8bd..1152ff6c2c27 100644 --- a/tests/ut/quantization/methods/test_w8a8_mxfp8.py +++ b/tests/ut/quantization/methods/test_w8a8_mxfp8.py @@ -110,7 +110,7 @@ def test_get_dynamic_quant_param_dtype_uint8(self): 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, [8, 512, 4]) + self.assertEqual(result["w13_weight_scale"].shape, (8, 512, 4)) self.assertEqual(result["w2_weight_scale"].dtype, torch.uint8) def test_process_weights_stores_original_shapes(self): diff --git a/tests/ut/quantization/test_method_adapters.py b/tests/ut/quantization/test_method_adapters.py index 263fb9ca08aa..95ada727bb06 100644 --- a/tests/ut/quantization/test_method_adapters.py +++ b/tests/ut/quantization/test_method_adapters.py @@ -15,7 +15,7 @@ class TestAscendLinearMethod(TestBase): @patch("vllm_ascend.quantization.method_adapters.enable_dsa_cp_with_layer_shard") - def setUp(self): + def setUp(self, mock_enable_dsa_cp_with_layer_shard): self.mock_scheme = MagicMock(spec=AscendLinearScheme) self.mock_scheme.get_weight.return_value = { "weight": torch.empty(128, 256, dtype=torch.int8),