[CI] fix

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

Author: zzhx1

tests/ut/models/test_deepseek_v2.py

@@ -22,8 +22,8 @@
 from vllm.distributed.parallel_state import GroupCoordinator
 
 from vllm_ascend.models.deepseek_v2 import (
-    CustomDeepseekV2ForCausalLM, CustomDeepseekV2MergedReplicatedLinear,
-    CustomDeepseekV2MLAAttention, CustomDeepseekV2MLP, CustomDeepseekV2MoE,
+    CustomDeepseekV2MergedReplicatedLinear, CustomDeepseekV2MLAAttention,
+    CustomDeepseekV2MLP, CustomDeepseekV2MoE,
     CustomDeepseekV2RowParallelLinear,
     CustomDeepseekV2RowParallelLinearReplaceAllreduce,
     CustomDeepseekV2SiluAndMul, LogitsProcessor, ParallelLMHead)
@@ -267,33 +267,29 @@ def test_custom_deepseek_v2_mla_attention(mock_rms_norm, mock_distributed,
                                         prefix="layers.1.self_attn")
     assert hasattr(attn, "q_proj")
 
+
 def test_deepseek_v2_lmhead(mock_distributed, vllm_config):
     # 创建一个简单的配置对象
     class SimpleConfig:
+
         def __init__(self):
             self.vocab_size = 10000
             self.hidden_size = 128
 
     config = SimpleConfig()
-    
+
     # 直接创建lmhead和logits_processor
     lmhead = ParallelLMHead(config.vocab_size, config.hidden_size)
     logits_processor = LogitsProcessor(config.vocab_size)
 
-    # 创建测试输入
-    input_ids = torch.randint(0, config.vocab_size, (2, 4))
-    positions = torch.arange(4).repeat(2, 1)
-
     # 创建模拟输出
     mock_output = torch.randn(2, 4, config.hidden_size)
     mock_logits = torch.randn(2, 4, config.vocab_size)
 
     # 直接测试logits_processor
-    with patch.object(lmhead.quant_method,
-                      "apply",
-                      return_value=mock_logits):
+    with patch.object(lmhead.quant_method, "apply", return_value=mock_logits):
         with patch.object(logits_processor,
                           "_gather_logits",
                           return_value=mock_logits):
             logits = logits_processor(lmhead, mock_output)
-    assert logits.shape == (2, 4, config.vocab_size)
+    assert logits.shape == (2, 4, config.vocab_size)

tests/ut/ops/test_vocab_parallel_embedding.py

@@ -18,8 +18,8 @@
 
 import torch
 
-from vllm_ascend.ops.vocab_parallel_embedding import \
-    AscendVocabParallelEmbedding
+from vllm_ascend.ops.vocab_parallel_embedding import (
+    AscendLogitsProcessor, AscendParallelLMHead, AscendVocabParallelEmbedding)
 
 VOCAB_PARALLEL_EMBEDDING_TEST_NUM_RANDOM_SEEDS = 128
 
@@ -178,3 +178,55 @@ def test_output_shape(self):
                     # Call the forward method
                     output = layer.forward(input_)
                 self.assertEqual(output.shape, expected_shape)
+
+
+class TestAscendLogitsProcessor(unittest.TestCase):
+
+    def setUp(self):
+        self.vocab_size = 50
+        self.num_embeddings = 50
+        self.embedding_dim = 10
+        self.org_num_embeddings = 40
+        self.padding_size = 8
+
+        self.mock_group = MagicMock()
+        self.mock_group.world_size = 2
+        self.mock_group.rank_in_group = 0
+        self.mock_ascend_config = MagicMock()
+        self.mock_quant_method = MagicMock()
+        self.mock_quant_method.apply = MagicMock(
+            return_value=torch.randn(1, self.vocab_size))
+        self.patches = [
+            patch("vllm_ascend.ascend_config.get_ascend_config",
+                  return_value=self.mock_ascend_config),
+            patch(
+                "vllm_ascend.ops.vocab_parallel_embedding.get_lmhead_tp_group",
+                return_value=self.mock_group),
+            patch("vllm_ascend.ops.vocab_parallel_embedding.lmhead_tp_enable",
+                  return_value=True),
+            patch(
+                "vllm_ascend.ops.vocab_parallel_embedding.get_lmhead_tp_group.all_to_all",
+                return_value=torch.randn(1, self.vocab_size))
+        ]
+
+        for p in self.patches:
+            p.start()
+
+    def tearDown(self):
+        for p in self.patches:
+            p.stop()
+
+    def test_create_processor(self):
+        processor = AscendLogitsProcessor(vocab_size=self.vocab_size)
+        self.assertEqual(processor.vocab_size, self.vocab_size)
+
+    def test_get_logits(self):
+        processor = AscendLogitsProcessor(vocab_size=self.vocab_size)
+        lmhead = AscendParallelLMHead(num_embeddings=self.num_embeddings,
+                                      embedding_dim=self.embedding_dim,
+                                      prefix="lm_head")
+        lmhead.quant_method = self.mock_quant_method
+        lmhead.quant_method.apply = self.mock_quant_method.apply
+        hidden_state = torch.randn(1, self.org_num_embeddings)
+        processor._get_logits(hidden_state, lmhead)
+        self.mock_quant_method.apply.assert_called_once()

vllm_ascend/ops/vocab_parallel_embedding.py

@@ -207,7 +207,7 @@ def _get_logits(
         self,
         hidden_states: torch.Tensor,
         lm_head: AscendParallelLMHead,
-        embedding_bias: Optional[torch.Tensor],
+        embedding_bias: Optional[torch.Tensor] = None,
     ) -> Optional[torch.Tensor]:
         if lmhead_tp_enable():
             return self._get_logits_lmheadtp(hidden_states, lm_head,

vllm_ascend/worker/model_runner_v1.py

@@ -1278,7 +1278,7 @@ def _prepare_inputs(
             logits_indices = spec_decode_metadata.logits_indices
 
         if lmhead_tp_enable():
-            max_num_reqs_across_dp = padded_num_tokens_across_dp if not with_prefill else self.max_num_reqs
+            max_num_reqs_across_dp = maybe_padded_num_tokens if not with_prefill else self.max_num_reqs
             logits_indices = nn.functional.pad(
                 logits_indices,
                 (0, max_num_reqs_across_dp - logits_indices.shape[0]))