Add ut aboult mtp fullgraph

Signed-off-by: anon189Ty <[email protected]>

Author: anon189Ty

tests/ut/attention/test_mla_v1.py

@@ -229,6 +229,41 @@ def test_ascend_mla_metadata_builder_spec_decode(self):
                 builder.chunked_prefill_enabled,
                 mock_vllm_config.scheduler_config.chunked_prefill_enabled)
 
+    def test_ascend_mla_metadata_builder_build_full_graph(self):
+        mock_vllm_config = MagicMock()
+        mock_vllm_config.model_config.max_model_len = 1024
+        mock_vllm_config.model_config.get_head_size.return_value = 64
+        mock_vllm_config.model_config.dtype = torch.float16
+        mock_vllm_config.cache_config.block_size = 16
+        mock_vllm_config.scheduler_config.max_num_seqs = 4
+        mock_vllm_config.scheduler_config.chunked_prefill_enabled = False
+        mock_device = 'cpu'
+
+        mock_spec_config = MagicMock()
+        mock_spec_config.num_speculative_tokens = 3
+        mock_vllm_config.speculative_config = mock_spec_config
+
+        builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config,
+                                           mock_device)
+        common_metadata = MagicMock()
+        model = MagicMock()
+        common_metadata.graph_pad_size = 8
+        common_metadata.num_reqs = 4
+        common_metadata.num_actual_tokens = 5
+        common_metadata.max_query_len = 5
+        common_metadata.seq_lens_cpu = torch.Tensor([9, 10, 8, 8]).int()
+        common_metadata.query_start_loc = torch.Tensor([0, 1, 2, 4, 5]).int()
+        common_metadata.query_start_loc_cpu = torch.Tensor([0, 1, 2, 4,
+                                                            5]).int()
+        common_metadata.positions = torch.Tensor([1, 2, 3, 4, 5, 6]).int()
+        block_table = torch.Tensor([[1, 0], [2, 0], [3, 0], [4, 0]]).int()
+        common_metadata.block_table_tensor = block_table
+        metadata = builder.build(0, common_metadata, model)
+
+        self.assertEqual(metadata.decode.actual_seq_lengths_q,
+                         [1, 2, 4, 5, 6, 6, 7, 8])
+        self.assertEqual(metadata.decode.block_table.shape[0], 8)
+
     def test_reorder_batch(self):
         ascend_config = MagicMock()
 
@@ -266,6 +301,28 @@ def test_reorder_batch(self):
         self.assertTrue(modified)
         input_batch.swap_states.assert_called_once_with(1, 2)
 
+    def test_pad_actual_seq_lens_q(self):
+        mock_vllm_config = MagicMock()
+        mock_vllm_config.model_config.max_model_len = 1024
+        mock_vllm_config.model_config.get_head_size.return_value = 64
+        mock_vllm_config.model_config.dtype = torch.float16
+        mock_vllm_config.cache_config.block_size = 16
+        mock_vllm_config.scheduler_config.max_num_seqs = 4
+        mock_vllm_config.scheduler_config.chunked_prefill_enabled = False
+        mock_device = 'cpu'
+        mock_vllm_config.speculative_config = None
+
+        builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config,
+                                           mock_device)
+        input_seq_lens = [1, 2, 4, 5]
+        expect_output = [1, 2, 4, 5, 6, 6, 7, 8]
+        num_reqs = 4
+        num_reqs_pad_size = 4
+        output_seq_lens = builder.pad_actual_seq_len_q(num_reqs_pad_size,
+                                                       num_reqs,
+                                                       input_seq_lens)
+        self.assertEqual(output_seq_lens, expect_output)
+
 
 class TestAscendMLAImpl(TestBase):
 

tests/ut/compilation/test_acl_graph.py

@@ -21,7 +21,9 @@
 from vllm.forward_context import BatchDescriptor, ForwardContext
 
 from tests.ut.base import TestBase
-from vllm_ascend.compilation.acl_graph import ACLGraphEntry, ACLGraphWrapper
+from vllm_ascend.compilation.acl_graph import (
+    ACLGraphEntry, ACLGraphWrapper, get_mtp_graph_params, set_mtp_graph_params,
+    update_mtp_graph_params_workspaces)
 
 
 class TestACLGraphEntry(TestBase):
@@ -718,3 +720,24 @@ def test_unwrap_method(self):
 
         unwrapped = wrapper.unwrap()
         self.assertEqual(unwrapped, self.mock_runnable)
+
+
+class TestMTPGraphParams(TestBase):
+
+    def test_set_mtp_graph_params(self):
+        with patch('vllm_ascend.compilation.acl_graph._mtp_graph_params',
+                   new=None):
+            set_mtp_graph_params([4])
+            from vllm_ascend.compilation.acl_graph import _mtp_graph_params
+            self.assertIsNotNone(_mtp_graph_params)
+
+    @patch('vllm_ascend.compilation.acl_graph._mtp_graph_params')
+    def test_update_mtp_graph_params_workspaces(self, mtp_graph_params_mock):
+        mtp_graph_params_mock.workspaces = {4: 5}
+        update_mtp_graph_params_workspaces(4, 6)
+        self.assertEqual(mtp_graph_params_mock.workspaces[4], 6)
+
+    @patch('vllm_ascend.compilation.acl_graph._mtp_graph_params')
+    def test_get_mtp_graph_params(self, mtp_graph_params_mock):
+        graph_params = get_mtp_graph_params()
+        self.assertIs(mtp_graph_params_mock, graph_params)

vllm_ascend/attention/mla_v1.py

@@ -451,8 +451,8 @@ def build(
                 num_reqs_pad_size = graph_pad_size - num_reqs
                 actual_seq_lengths_q = self.pad_actual_seq_len_q(
                     num_reqs_pad_size, num_reqs, actual_seq_lengths_q)
-                seq_lens_list = seq_lens_list + [0] * (
-                    graph_pad_size - num_decodes)
+                seq_lens_list = seq_lens_list + [0] * (graph_pad_size -
+                                                       num_decodes)
                 num_block_pad_size = graph_pad_size - block_table.shape[0]
                 if num_block_pad_size > 0:
                     block_table_padding = torch.zeros(

vllm_ascend/spec_decode/mtp_proposer.py

@@ -36,6 +36,9 @@
 
 class MtpProposer(Proposer):
 
+    # TODO: Find out why ModuleRunner does not this explicit typing?
+    model: Union[nn.Module, ACLGraphWrapper]
+
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -145,7 +148,8 @@ def dummy_run(self,
         if skip_attn:
             attn_metadata = None
         elif is_running_torchair:
-            common_attn_metadata = TorchairCommonAttentionMetadata(
+            common_attn_metadata: TorchairCommonAttentionMetadata = \
+            TorchairCommonAttentionMetadata(
                 num_reqs=num_reqs,
                 num_actual_tokens=1,
                 actual_seq_lengths_q=self.runner.actual_seq_lengths_q,
@@ -156,24 +160,18 @@ def dummy_run(self,
             attn_metadata = self.runner.attn_metadata_builder.build_torchair_graph_dummy(
                 common_attn_metadata)
         elif aclgraph_runtime_mode == CUDAGraphMode.FULL:
-            # assert with_prefill is False, \
-            #     "Full decode graph only supports uniform batch now."
-            max_seq_lens = self.runner.model_config.max_model_len
             if len(self.runner.attn_groups) > 0:
                 num_computed_tokens_cpu = (
                     self.runner.input_batch.
                     num_computed_tokens_cpu_tensor[:num_reqs])
-                query_start_loc = torch.tensor(
-                    [0] + self.runner.actual_seq_lengths_q[:num_reqs],
-                    device=self.runner.device,
-                    dtype=torch.int32)
-                common_attn_metadata = AscendCommonAttentionMetadata(
+                common_attn_metadata: AscendCommonAttentionMetadata = \
+                AscendCommonAttentionMetadata(
                     query_start_loc=torch.tensor(
                         [0] + self.runner.actual_seq_lengths_q[:num_reqs],
                         device=self.device,
                         dtype=torch.int32),
-                    query_start_loc_cpu=self.runner.query_start_loc_cpu[
-                        :num_reqs + 1],
+                    query_start_loc_cpu=self.runner.
+                    query_start_loc_cpu[:num_reqs + 1],
                     seq_lens_cpu=self.runner.seq_lens_cpu,
                     seq_lens=self.runner.seq_lens_cpu[:num_reqs],
                     num_reqs=num_reqs,
@@ -183,7 +181,8 @@ def dummy_run(self,
                     actual_seq_lengths_q=self.runner.actual_seq_lengths_q,
                     block_table_tensor=self.runner.input_batch.block_table[0].
                     get_device_tensor()[:num_reqs],
-                    slot_mapping=self.runner.input_batch.block_table[0].slot_mapping,
+                    slot_mapping=self.runner.input_batch.block_table[0].
+                    slot_mapping,
                     positions=self.runner.positions,
                     attn_mask=self.runner.attn_mask,
                     spec_attn_mask=self.runner.spec_attn_mask,
@@ -466,7 +465,6 @@ def _propose(
 
         seq_lens = target_positions[last_token_indices] + 1
         seq_lens = seq_lens.int()
-        seq_lens_len = seq_lens.shape[0]
 
         if not self.torchair_graph_enabled:
             # torch mode need to update num_tokens_across_dp
@@ -481,10 +479,10 @@ def _propose(
 
         if scheduler_output:
             uniform_decode = (max_query_len in list(
-                range(1, self.num_speculative_tokens + 2))) and (
-                    scheduler_output.total_num_scheduled_tokens ==
-                    self.runner.input_batch.num_reqs *
-                    (self.num_speculative_tokens + 1))
+                range(1, self.num_speculative_tokens +
+                      2))) and (scheduler_output.total_num_scheduled_tokens
+                                == self.runner.input_batch.num_reqs *
+                                (self.num_speculative_tokens + 1))
             batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens,
                                                uniform_decode=uniform_decode)
         else:
@@ -500,9 +498,12 @@ def _propose(
             # Currently, if not torchair, runner.graph_pad_size will always be -1.
             graph_pad_size = self.runner.graph_pad_size
 
-        runner_slot_mapping = self.runner.input_batch.block_table[0].slot_mapping
-        runner_slot_mapping[:target_slot_mapping.shape[0]].copy_(target_slot_mapping)
-        runner_slot_mapping[target_slot_mapping.shape[0]:num_input_tokens].fill_(0)
+        runner_slot_mapping = self.runner.input_batch.block_table[
+            0].slot_mapping
+        runner_slot_mapping[:target_slot_mapping.shape[0]].copy_(
+            target_slot_mapping)
+        runner_slot_mapping[target_slot_mapping.
+                            shape[0]:num_input_tokens].fill_(0)
 
         # NOTE: Currently, just positions, slot_mapping, block_table and
         # seq_lens will be sent into MLAMetadata.