add patch ut

Pr0Wh1teGivee · Pr0Wh1teGivee · commit 36fe57f76c33 · 2025-07-21T17:43:57.000+08:00
Signed-off-by: Pr0Wh1teGivee &lt;calvin_zhu0210@outlook.com&gt;
diff --git a/tests/ut/patch/worker/patch_common/test_patch_distributed.py b/tests/ut/patch/worker/patch_common/test_patch_distributed.py
@@ -11,17 +11,106 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
-#
+
+from unittest.mock import MagicMock, patch
+
+import torch
+from vllm.distributed.parallel_state import GroupCoordinator
 
 from tests.ut.base import TestBase
+from vllm_ascend.patch.worker.patch_common.patch_distributed import \
+    GroupCoordinatorPatch
 
 
-class TestPatchDistributed(TestBase):
+class TestGroupCoordinatorPatch(TestBase):
 
-    def test_GroupCoordinator_patched(self):
-        from vllm.distributed.parallel_state import GroupCoordinator
+    def setUp(self):
+        self.mock_group_ranks = [[0, 1]]
+        self.mock_local_rank = 0
+        self.mock_backend = "hccl"
+        self.mock_use_device_comm = True
+
+        patcher_get_rank = patch("torch.distributed.get_rank", return_value=0)
+        patcher_new_group = patch("torch.distributed.new_group",
+                                  return_value=MagicMock())
+        patcher_is_cuda_alike = patch(
+            "vllm.platforms.current_platform.is_cuda_alike", return_value=True)
+        patcher_device_comm_cls = patch(
+            "vllm.distributed.parallel_state.resolve_obj_by_qualname",
+            return_value=MagicMock())
 
-        from vllm_ascend.patch.worker.patch_common.patch_distributed import \
-            GroupCoordinatorPatch
+        self.mock_get_rank = patcher_get_rank.start()
+        self.mock_new_group = patcher_new_group.start()
+        self.mock_is_cuda_alike = patcher_is_cuda_alike.start()
+        self.mock_resolve_obj = patcher_device_comm_cls.start()
 
+        self.addCleanup(patcher_get_rank.stop)
+        self.addCleanup(patcher_new_group.stop)
+        self.addCleanup(patcher_is_cuda_alike.stop)
+        self.addCleanup(patcher_device_comm_cls.stop)
+
+        self.group_coordinator = GroupCoordinatorPatch(
+            group_ranks=self.mock_group_ranks,
+            local_rank=self.mock_local_rank,
+            torch_distributed_backend=self.mock_backend,
+            use_device_communicator=self.mock_use_device_comm)
+
+    def test_GroupCoordinator_patched(self):
         self.assertIs(GroupCoordinator, GroupCoordinatorPatch)
+
+    def test_initialization_sets_attributes(self):
+        self.assertEqual(self.group_coordinator.world_size, 2)
+        self.assertEqual(self.group_coordinator.rank_in_group, 0)
+        self.assertTrue(hasattr(self.group_coordinator, "device_communicator"))
+
+    def test_all_to_all_returns_input_when_world_size_1(self):
+        self.group_coordinator.world_size = 1
+        input_tensor = torch.randn(2, 3)
+        output = self.group_coordinator.all_to_all(input_tensor)
+        self.assertTrue(torch.equal(output, input_tensor))
+
+    def test_all_to_all_raises_assertion_on_invalid_scatter_dim(self):
+        input_tensor = torch.randn(2, 3)
+        with self.assertRaises(AssertionError) as cm:
+            self.group_coordinator.all_to_all(input_tensor, scatter_dim=2)
+        self.assertIn("Invalid scatter dim", str(cm.exception))
+
+    def test_all_to_all_raises_assertion_on_invalid_gather_dim(self):
+        input_tensor = torch.randn(2, 3)
+        with self.assertRaises(AssertionError) as cm:
+            self.group_coordinator.all_to_all(input_tensor, gather_dim=2)
+        self.assertIn("Invalid gather dim", str(cm.exception))
+
+    def test_all_to_all_calls_device_communicator_with_correct_args(self):
+        mock_communicator = MagicMock()
+        self.group_coordinator.device_communicator = mock_communicator
+
+        input_tensor = torch.randn(2, 3)
+        scatter_dim = 0
+        gather_dim = 1
+        scatter_sizes = [1, 1]
+        gather_sizes = [1, 1]
+
+        self.group_coordinator.all_to_all(input_tensor,
+                                          scatter_dim=scatter_dim,
+                                          gather_dim=gather_dim,
+                                          scatter_sizes=scatter_sizes,
+                                          gather_sizes=gather_sizes)
+
+        mock_communicator.all_to_all.assert_called_once_with(
+            input_tensor, scatter_dim, gather_dim, scatter_sizes, gather_sizes)
+
+    def test_all_to_all_calls_device_communicator_without_sizes(self):
+        mock_communicator = MagicMock()
+        self.group_coordinator.device_communicator = mock_communicator
+
+        input_tensor = torch.randn(2, 3)
+        scatter_dim = 0
+        gather_dim = 1
+
+        self.group_coordinator.all_to_all(input_tensor,
+                                          scatter_dim=scatter_dim,
+                                          gather_dim=gather_dim)
+
+        mock_communicator.all_to_all.assert_called_once_with(
+            input_tensor, scatter_dim, gather_dim, None, None)
diff --git a/tests/ut/patch/worker/patch_common/test_patch_minicpm.py b/tests/ut/patch/worker/patch_common/test_patch_minicpm.py
@@ -0,0 +1,76 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+
+from unittest.mock import MagicMock
+
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend.patch.worker.patch_common.patch_minicpm import forward
+
+
+class TestPatchedMiniCPMForward(TestBase):
+
+    def setUp(self):
+        self.mock_self = MagicMock()
+
+        self.mock_self.q_size = 128
+        self.mock_self.kv_size = 128
+
+        self.mock_self.qkv_proj = MagicMock()
+        self.mock_self.rotary_emb = MagicMock()
+        self.mock_self.attn = MagicMock()
+        self.mock_self.o_proj = MagicMock()
+
+        self.positions = torch.tensor([1, 2, 3])
+        self.hidden_states = torch.randn(3, 256)  # [batch_size, hidden_size]
+
+        self.mock_qkv = torch.randn(3, 384)
+        self.mock_q = self.mock_qkv[:, :128]
+        self.mock_k = self.mock_qkv[:, 128:256]
+        self.mock_v = self.mock_qkv[:, 256:]
+
+        self.mock_self.qkv_proj.return_value = (self.mock_qkv, None)
+        self.mock_self.rotary_emb.return_value = (self.mock_q, self.mock_k)
+        self.mock_self.attn.return_value = torch.randn(3, 256)
+        self.mock_self.o_proj.return_value = (torch.randn(3, 256), None)
+
+    def test_forward_patched(self):
+        from vllm.model_executor.models.minicpm import MiniCPMAttention
+
+        self.assertIs(MiniCPMAttention.forward, forward)
+
+    def test_forward_function(self):
+        result = forward(self.mock_self, self.positions, self.hidden_states)
+
+        self.mock_self.qkv_proj.assert_called_once_with(self.hidden_states)
+
+        args, _ = self.mock_self.rotary_emb.call_args
+        self.assertEqual(len(args), 3)
+        self.assertTrue(torch.equal(args[0], self.positions))
+        self.assertTrue(torch.equal(args[1], self.mock_q))
+        self.assertTrue(torch.equal(args[2], self.mock_k))
+
+        args, _ = self.mock_self.attn.call_args
+        self.assertEqual(len(args), 3)
+        self.assertTrue(torch.equal(args[0], self.mock_q))
+        self.assertTrue(torch.equal(args[1], self.mock_k))
+        self.assertTrue(torch.equal(args[2], self.mock_v))
+
+        self.mock_self.o_proj.assert_called_once_with(
+            self.mock_self.attn.return_value)
+
+        self.assertEqual(result.shape, (3, 256))
+        self.assertTrue(
+            torch.equal(result, self.mock_self.o_proj.return_value[0]))
diff --git a/tests/ut/patch/worker/patch_common/test_patch_multistep_worker.py b/tests/ut/patch/worker/patch_common/test_patch_multistep_worker.py
@@ -0,0 +1,154 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+
+from unittest.mock import MagicMock, patch
+
+import torch
+from vllm.model_executor.layers.sampler import SamplerOutput
+from vllm.sequence import ExecuteModelRequest, SequenceGroupMetadata
+
+from tests.ut.base import TestBase
+from vllm_ascend.patch.worker.patch_common.patch_multi_step_worker import \
+    sampler_output
+
+
+class TestPatchedMultiStepWorkerSamplerOutput(TestBase):
+
+    def setUp(self):
+        self.mock_self = MagicMock()
+
+        self.mock_self.device = torch.device("cpu")
+
+        self.mock_self._raise_if_unsupported = MagicMock()
+        self.mock_self._expand_execute_model_request = MagicMock()
+        self.mock_self.execute_model = MagicMock()
+        self.mock_self._maybe_update_previous_hidden_states = MagicMock()
+        self.mock_self._append_new_tokens = MagicMock()
+        self.mock_self._filter_model_output = MagicMock()
+
+        self.execute_model_req = ExecuteModelRequest(
+            seq_group_metadata_list=[MagicMock(spec=SequenceGroupMetadata)],
+            num_steps=1,
+            blocks_to_swap_in={},
+            blocks_to_swap_out={},
+            blocks_to_copy={},
+            num_lookahead_slots=0)
+        self.sample_len = 3
+        self.seq_ids_with_bonus_token = {1, 2, 3}
+
+        self.expanded_request = MagicMock(spec=ExecuteModelRequest)
+        self.indices_of_seq_with_bonus_tokens = [0, 1, 2]
+        self.mock_self._expand_execute_model_request.return_value = (
+            self.expanded_request, self.indices_of_seq_with_bonus_tokens)
+
+        self.filtered_output = [
+            MagicMock(spec=SamplerOutput),
+            MagicMock(spec=SamplerOutput),
+            MagicMock(spec=SamplerOutput)
+        ]
+        self.mock_self._filter_model_output.return_value = self.filtered_output
+
+    def test_sampler_output_patched(self):
+        from vllm.spec_decode.multi_step_worker import MultiStepWorker
+
+        wrapped_func = MultiStepWorker.sampler_output.__wrapped__
+        self.assertIs(
+            wrapped_func, sampler_output,
+            "Wrapped function does not match the expected implementation")
+
+    def test_gpu_multi_step_path(self):
+        mock_model_runner = MagicMock()
+        mock_model_runner.supports_gpu_multi_step.return_value = True
+
+        self.mock_self.model_runner = mock_model_runner
+        with patch(
+                'vllm_ascend.patch.worker.patch_common.patch_multi_step_worker.isinstance'
+        ) as mock_isinstance:
+            mock_isinstance.return_value = True
+
+            mock_outputs = [
+                MagicMock(spec=SamplerOutput),
+                MagicMock(spec=SamplerOutput),
+                MagicMock(spec=SamplerOutput)
+            ]
+            self.mock_self.execute_model.return_value = mock_outputs
+
+            result, need_transpose = sampler_output(
+                self.mock_self, self.execute_model_req, self.sample_len,
+                self.seq_ids_with_bonus_token)
+
+        self.mock_self._raise_if_unsupported.assert_called_once_with(
+            self.execute_model_req)
+        self.mock_self._expand_execute_model_request.assert_called_once_with(
+            self.execute_model_req, self.seq_ids_with_bonus_token)
+
+        mock_model_runner.supports_gpu_multi_step.assert_called_once_with(
+            self.expanded_request)
+        self.assertEqual(self.expanded_request.num_steps, self.sample_len)
+        mock_model_runner.set_indices_of_seq_with_bonus_tokens.assert_called_once_with(
+            self.indices_of_seq_with_bonus_tokens)
+        self.mock_self.execute_model.assert_called_once_with(
+            execute_model_req=self.expanded_request)
+
+        self.assertEqual(result, self.filtered_output)
+        self.assertTrue(need_transpose)
+
+        self.mock_self._maybe_update_previous_hidden_states.assert_not_called()
+        self.mock_self._append_new_tokens.assert_not_called()
+
+    def test_cpu_multi_step_path(self):
+        mock_model_runner = MagicMock()
+        mock_model_runner.supports_gpu_multi_step.return_value = False
+
+        self.mock_self.model_runner = mock_model_runner
+        self.mock_self.worker = MagicMock()
+
+        mock_step_output = MagicMock(spec=SamplerOutput)
+        self.mock_self.worker.execute_model.return_value = [[mock_step_output]]
+
+        result, need_transpose = sampler_output(self.mock_self,
+                                                self.execute_model_req,
+                                                self.sample_len,
+                                                self.seq_ids_with_bonus_token)
+
+        self.assertEqual(self.mock_self.worker.execute_model.call_count,
+                         self.sample_len)
+        self.mock_self._append_new_tokens.assert_called()
+        self.assertEqual(self.mock_self._append_new_tokens.call_count,
+                         self.sample_len)
+
+        self.mock_self._filter_model_output.assert_called_once()
+        self.assertEqual(result, self.filtered_output)
+        self.assertTrue(need_transpose)
+
+    def test_cpu_path_with_hidden_states(self):
+        self.expanded_request.previous_hidden_states = MagicMock()
+
+        mock_model_runner = MagicMock()
+        mock_model_runner.supports_gpu_multi_step.return_value = False
+        self.mock_self.model_runner = mock_model_runner
+        self.mock_self.worker = MagicMock()
+
+        self.mock_self.worker.model_runner = MagicMock()
+        self.mock_self.worker.model_runner.return_hidden_states = False
+
+        mock_step_output = MagicMock(spec=SamplerOutput)
+        self.mock_self.worker.execute_model.return_value = [[mock_step_output]]
+
+        sampler_output(self.mock_self, self.execute_model_req, self.sample_len,
+                       self.seq_ids_with_bonus_token)
+
+        self.assertTrue(
+            self.mock_self.worker.model_runner.return_hidden_states)
+        self.mock_self._maybe_update_previous_hidden_states.assert_called()
diff --git a/tests/ut/patch/worker/patch_common/test_patch_utils.py b/tests/ut/patch/worker/patch_common/test_patch_utils.py
