fix rope ut

Signed-off-by: MengqingCao <cmq0113@163.com>

Author: MengqingCao

tests/ut/ops/test_rotary_embedding.py

@@ -1,15 +1,13 @@
 import math
 import unittest
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, PropertyMock, patch
 
 import torch
+from vllm.model_executor.layers.rotary_embedding import (
+    DeepseekScalingRotaryEmbedding, RotaryEmbedding)
 
 from tests.ut.base import TestBase
-from vllm_ascend.ops.rotary_embedding import (custom_rotary_embedding_enabled,
-                                              native_rope_deepseek_forward,
-                                              rope_forward_oot, rotate_half,
-                                              yarn_find_correction_dim,
-                                              yarn_get_mscale)
+from vllm_ascend.ops.rotary_embedding import custom_rotary_embedding_enabled
 
 
 class TestCustomRotaryEmbeddingEnabled(unittest.TestCase):
@@ -67,22 +65,28 @@ def test_custom_rotary_embedding_enabled(self):
             self.assertFalse(result)
 
 
-class TestRopeForwardOot(unittest.TestCase):
+class TestAscendRotaryEmbedding(unittest.TestCase):
 
     def setUp(self):
         # Common setup for tests
         self.positions = torch.tensor([1, 2, 3])
-        self.query = torch.randn(3, 4, dtype=torch.float16)
-        self.key = torch.randn(3, 4, dtype=torch.float16)
+        self.query = torch.randn(3, 1, 32, dtype=torch.float16)
+        self.key = torch.randn(3, 1, 32, dtype=torch.float16)
         self.head_size = 32
-        self.cos_sin_cache = torch.randn(3, 4)
+        self.rotary_dim = self.head_size
+        self.max_position = 16
+        self.rope_theta = 10000
+        self.is_neox_style = True
+        self.cos_sin_cache = torch.randn(3, 1, 32)
+        self.layer = RotaryEmbedding(self.head_size, self.rotary_dim,
+                                     self.max_position, self.rope_theta,
+                                     self.is_neox_style, torch.float16)
 
         # Mock self object for rope_forward_oot
         self.mock_self = MagicMock()
         self.mock_self.head_size = self.head_size
         self.mock_self.cos_sin_cache = self.cos_sin_cache
-        self.mock_self.is_neox_style = True
-        self.mock_self.forward_native.return_value = (self.query, self.key)
+        self.mock_self.is_neox_style = self.is_neox_style
 
     @patch('vllm_ascend.ops.rotary_embedding.get_ascend_config')
     def test_rope_forward_oot_torchair_enabled_base(self,
@@ -91,12 +95,14 @@ def test_rope_forward_oot_torchair_enabled_base(self,
         mock_config = MagicMock()
         mock_config.torchair_graph_config.enabled = True
         mock_get_ascend_config.return_value = mock_config
-
-        result_q, result_k = rope_forward_oot(self.mock_self, self.positions,
-                                              self.query, self.key)
-
-        self.mock_self.forward_native.assert_called_once_with(
-            self.positions, self.query, self.key, None)
+        with patch.object(self.layer,
+                          "forward_native",
+                          return_value=(self.query,
+                                        self.key)) as mock_forward_native:
+            result_q, result_k = self.layer.forward(self.positions, self.query,
+                                                    self.key)
+
+        mock_forward_native.assert_called_once()
         self.assertTrue(torch.equal(result_q, self.query))
         self.assertTrue(torch.equal(result_k, self.key))
 
@@ -117,9 +123,10 @@ def test_rope_forward_oot_custom_kernel(self, mock_rotary_embedding,
 
         mock__c.rotary_embedding.return_value = self.query, self.key
 
-        result_q, result_k = rope_forward_oot(self.mock_self, self.positions,
-                                              self.query, self.key)
+        result_q, result_k = self.layer.forward(self.positions, self.query,
+                                                self.key)
 
+        mock__c.rotary_embedding.assert_called_once()
         self.assertEqual(result_q.shape, self.query.shape)
         self.assertEqual(result_k.shape, self.key.shape)
 
@@ -138,8 +145,9 @@ def test_rope_forward_oot_contiguous(self, mock_npu_rotary,
         non_contig_query = self.query.transpose(0, 1)
         non_contig_key = self.key.transpose(0, 1)
 
-        result_q, result_k = rope_forward_oot(self.mock_self, self.positions,
-                                              non_contig_query, non_contig_key)
+        result_q, result_k = self.layer.forward(self.positions,
+                                                non_contig_query,
+                                                non_contig_key)
 
         mock_npu_rotary.assert_called_once()
         self.assertEqual(result_q.shape, non_contig_query.shape)
@@ -154,8 +162,7 @@ def test_rope_forward_oot_with_offsets(self, mock_get_ascend_config):
         # Test that NotImplementedError is raised when offsets is provided
         offsets = torch.tensor([1, 2, 3])
         with self.assertRaises(NotImplementedError):
-            rope_forward_oot(self.mock_self, self.positions, self.query,
-                             self.key, offsets)
+            self.layer.forward(self.positions, self.query, self.key, offsets)
 
     @patch('vllm_ascend.ops.rotary_embedding.get_ascend_config')
     @patch('vllm_ascend.ops.rotary_embedding.custom_rotary_embedding_enabled',
@@ -169,11 +176,10 @@ def test_rope_forward_oot_neox_style_override(self, mock_npu_rotary,
         mock_get_ascend_config.return_value = mock_config
 
         # Test neox_style override
-        result_q, result_k = rope_forward_oot(self.mock_self,
-                                              self.positions,
-                                              self.query,
-                                              self.key,
-                                              is_neox_style_override=False)
+        result_q, result_k = self.layer.forward(self.positions,
+                                                self.query,
+                                                self.key,
+                                                is_neox_style_override=False)
 
         # Check that neox_style=False was passed to the NPU function
         args, kwargs = mock_npu_rotary.call_args
@@ -191,98 +197,108 @@ def __init__(self, max_seq_len=2048, is_neox_style=True):
         self.base = 1
 
 
-class TestNativeRopeDeepseekForward(TestBase):
+class TestAscendDeepseekScalingRotaryEmbedding(TestBase):
 
-    @patch('vllm_ascend.ops.rotary_embedding.rope_forward_oot')
-    def test_native_rope_deepseek_forward_base(self, mock_rope_forward_oot):
-        module = MockRopeModule()
-        positions = torch.tensor([1, 2, 3])
-        query = torch.randn(1, 8, 128)
-        key = torch.randn(1, 8, 128)
-
-        mock_rope_forward_oot.return_value = (query, key)
-
-        q_pe, k_pe = native_rope_deepseek_forward(module, positions, query,
-                                                  key)
-
-        assert q_pe.shape == query.shape
-        assert k_pe.shape == key.shape
+    def setUp(self):
+        # Common setup for tests
+        self.positions = torch.tensor([1, 2, 3])
+        self.query = torch.randn(3, 1, 32, dtype=torch.float16)
+        self.key = torch.randn(3, 1, 32, dtype=torch.float16)
+        self.head_size = 32
+        self.rotary_dim = self.head_size
+        self.max_position = 16
+        self.rope_theta = 10000
+        self.is_neox_style = True
+        self.scaling_factor = 1
+        self.layer = None
+
+    def _create_layer(self):
+        self.layer = DeepseekScalingRotaryEmbedding(
+            self.head_size, self.rotary_dim, self.max_position,
+            self.rope_theta, self.is_neox_style, self.scaling_factor,
+            torch.float16)
+        return self.layer
+
+    @patch("vllm.model_executor.layers.rotary_embedding.current_platform",
+           new_callable=PropertyMock)
+    def test_native_rope_deepseek_forward_base(self, mock_current_platform):
+        mock_current_platform.device_type = torch.device("cpu")
+        self.layer = self._create_layer()
+        with patch("vllm_ascend.ops.rotary_embedding.rope_forward_oot",
+                   return_value=(self.query,
+                                 self.key)) as mock_rope_forward_oot:
+            q_pe, k_pe = self.layer.forward(self.positions, self.query,
+                                            self.key)
+        mock_rope_forward_oot.assert_called_once()
+        assert q_pe.shape == self.query.shape
+        assert k_pe.shape == self.key.shape
 
-    @patch('vllm_ascend.ops.rotary_embedding._set_cos_sin_cache')
     @patch('vllm_ascend.ops.rotary_embedding.rope_forward_oot')
+    @patch("vllm.model_executor.layers.rotary_embedding.current_platform",
+           new_callable=PropertyMock)
     def test_native_rope_deepseek_forward_cache_handling(
-            self, mock_rope_forward_oot, mock_set_cache):
+            self, mock_current_platform, mock_rope_forward_oot):
+        mock_current_platform.device_type = torch.device("cpu")
+        self.layer = self._create_layer()
+        self.layer.max_seq_len = 1024
         # Test cache situation is true
-        module = MockRopeModule(max_seq_len=1024)
-        positions = torch.tensor([1, 2, 3])
-        query = torch.randn(1, 8, 128)
-        key = torch.randn(1, 8, 128)
+        with patch.object(self.layer, "_set_cos_sin_cache") as mock_set_cache:
+            mock_rope_forward_oot.return_value = (self.query, self.key)
 
-        mock_rope_forward_oot.return_value = (query, key)
-
-        q_pe, k_pe = native_rope_deepseek_forward(module,
-                                                  positions,
-                                                  query,
-                                                  key,
-                                                  max_seq_len=2048)
-
-        assert q_pe.shape == query.shape
-        assert k_pe.shape == key.shape
+            q_pe, k_pe = self.layer.forward(self.positions,
+                                            self.query,
+                                            self.key,
+                                            max_seq_len=2048)
+        mock_set_cache.assert_called_once()
+        assert q_pe.shape == self.query.shape
+        assert k_pe.shape == self.key.shape
 
     @patch('vllm_ascend.ops.rotary_embedding.rope_forward_oot')
+    @patch("vllm.model_executor.layers.rotary_embedding.current_platform",
+           new_callable=PropertyMock)
     def test_native_rope_deepseek_forward_key_reshaping(
-            self, mock_rope_forward_oot):
-        module = MockRopeModule()
-        positions = torch.tensor([1, 2, 3])
-        query = torch.randn(1, 8, 128)
-        key = torch.randn(1, 128)
+            self, mock_current_platform, mock_rope_forward_oot):
+        mock_current_platform.device_type = torch.device("cpu")
+        self.layer = self._create_layer()
 
-        mock_rope_forward_oot.return_value = (query, key)
+        key = torch.randn(1, 32)
 
-        q_pe, k_pe = native_rope_deepseek_forward(module, positions, query,
-                                                  key)
+        mock_rope_forward_oot.return_value = (self.query, key)
 
-        assert q_pe.shape == query.shape
-        assert k_pe.shape == (1, 128)
+        q_pe, k_pe = self.layer.forward(self.positions, self.query, key)
+        mock_rope_forward_oot.assert_called_once()
+        assert q_pe.shape == self.query.shape
+        assert k_pe.shape == key.shape
 
     @patch('vllm_ascend.ops.rotary_embedding.rope_forward_oot')
+    @patch("vllm.model_executor.layers.rotary_embedding.current_platform",
+           new_callable=PropertyMock)
     def test_native_rope_deepseek_forward_non_neox_style(
-            self, mock_rope_forward_oot):
-        module = MockRopeModule(is_neox_style=False)
-        positions = torch.tensor([1, 2, 3])
-        query = torch.randn(1, 8, 128)
-        key = torch.randn(1, 8, 128)
-
-        mock_rope_forward_oot.return_value = (query, key)
+            self, mock_current_platform, mock_rope_forward_oot):
+        mock_current_platform.device_type = torch.device("cpu")
+        self.layer = self._create_layer()
 
-        q_pe, k_pe = native_rope_deepseek_forward(module, positions, query,
-                                                  key)
-
-        assert q_pe.shape == query.shape
-        assert k_pe.shape == key.shape
+        mock_rope_forward_oot.return_value = (self.query, self.key)
 
+        q_pe, k_pe = self.layer.forward(self.positions, self.query, self.key)
 
-class TestRotateHalf(unittest.TestCase):
-
-    def test_rotate_half_even_dim(self):
-        # Test with even dimension
-        x = torch.tensor([1.0, 2.0, 3.0, 4.0])
-        expected = torch.tensor([-3.0, -4.0, 1.0, 2.0])
-        result = rotate_half(x)
-        self.assertTrue(torch.allclose(result, expected))
+        mock_rope_forward_oot.assert_called_once()
+        assert q_pe.shape == self.query.shape
+        assert k_pe.shape == self.key.shape
 
-
-class TestYarnFindCorrectionDim(unittest.TestCase):
-
-    def test_basic_case(self):
+    @patch("vllm.model_executor.layers.rotary_embedding.current_platform",
+           new_callable=PropertyMock)
+    def test_basic_case(self, mock_current_platform):
         # Test with standard values
+        mock_current_platform.device_type = torch.device("cpu")
+        self.layer = self._create_layer()
         num_rotations = 100
         dim = 512
         base = 10000
         max_position_embeddings = 2048
 
-        result = yarn_find_correction_dim(num_rotations, dim, base,
-                                          max_position_embeddings)
+        result = self.layer._yarn_find_correction_dim(num_rotations, dim, base,
+                                                      max_position_embeddings)
 
         # Calculate expected value manually
         expected = (dim * torch.log(
@@ -292,22 +308,25 @@ def test_basic_case(self):
 
         self.assertTrue(torch.allclose(result, expected))
 
+    @patch("vllm.model_executor.layers.rotary_embedding.current_platform",
+           new_callable=PropertyMock)
+    def test_yarn_get_mscale(self, mock_current_platform):
+        mock_current_platform.device_type = torch.device("cpu")
+        self.layer = self._create_layer()
 
-class TestYarnGetMscale(unittest.TestCase):
-
-    def test_scale_less_than_or_equal_1(self):
-        self.assertEqual(yarn_get_mscale(scale=0.5), 1.0)
-        self.assertEqual(yarn_get_mscale(scale=1.0), 1.0)
-        self.assertEqual(yarn_get_mscale(scale=0.999), 1.0)
+        # test_scale_less_than_or_equal_1
+        self.assertEqual(self.layer._yarn_get_mscale(scale=0.5), 1.0)
+        self.assertEqual(self.layer._yarn_get_mscale(scale=1.0), 1.0)
+        self.assertEqual(self.layer._yarn_get_mscale(scale=0.999), 1.0)
 
-    def test_scale_greater_than_1(self):
+        # test_scale_greater_than_1:
         test_cases = [(2.0, 1.0, 1.0 + 0.1 * math.log(2.0)),
                       (10.0, 1.0, 1.0 + 0.1 * math.log(10.0)),
                       (5.0, 2.0, 1.0 + 0.2 * math.log(5.0)),
                       (math.e, 1.0, 1.0 + 0.1)]
 
         for scale, mscale, expected in test_cases:
-            result = yarn_get_mscale(scale, mscale)
+            result = self.layer._yarn_get_mscale(scale, mscale)
             self.assertAlmostEqual(
                 result,
                 expected,