[1/N][UT][v1 MTP] add basic v1 mtp features

Co-authored-by: XWFAlone xuewenfei2@huawei.com
Co-authored-by: mengwei805 <mengwei25@huawei.com>
Co-authored-by: JC-ut0 xuyexiong@huawei.com
Signed-off-by: XWFAlone <xuewenfei2@huawei.com>

Author: XWFAlone

.github/workflows/vllm_ascend_test.yaml

@@ -163,7 +163,10 @@ jobs:
           if [[ "${{ matrix.os }}" == "linux-arm64-npu-1" ]]; then
             VLLM_USE_MODELSCOPE=true pytest -sv tests/singlecard/spec_decode/e2e/test_v1_spec_decode.py
             pytest -sv tests/singlecard/spec_decode/e2e/test_mtp_correctness.py  # it needs a clean process
-            pytest -sv tests/singlecard/spec_decode --ignore=tests/singlecard/spec_decode/e2e/test_mtp_correctness.py --ignore=tests/singlecard/spec_decode/e2e/test_v1_spec_decode.py
+            pytest -sv tests/singlecard/spec_decode --ignore=tests/singlecard/spec_decode/e2e/test_mtp_correctness.py --ignore=tests/singlecard/spec_decode/e2e/test_v1_spec_decode.py --ignore=tests/singlecard/spec_decode/e2e/test_v1_mtp_correctness.py
+          fi
+          if [[ "${{ matrix.os }}" == "linux-arm64-npu-1" ]] && [[ "${{ matrix.vllm_verison }}" == "main" ]]; then
+            VLLM_USE_MODELSCOPE=true pytest -sv tests/singlecard/spec_decode/e2e/test_v1_mtp_correctness.py
           fi
 
       - name: Run vllm-project/vllm test for V0 Engine

tests/singlecard/spec_decode/e2e/test_v1_mtp_correctness.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import os
+import random
+from typing import Any
+
+import pytest
+from vllm import LLM, SamplingParams
+
+os.environ['VLLM_USE_MODELSCOPE'] = 'True'
+
+
+@pytest.fixture
+def test_prompts():
+    prompt_types = ["repeat", "sentence"]
+    num_prompts = 10
+    prompts = []
+
+    random.seed(0)
+    random_prompt_type_choices = random.choices(prompt_types, k=num_prompts)
+
+    # Generate a mixed batch of prompts, some of which can be easily
+    # predicted by n-gram matching and some which likely cannot.
+    for kind in random_prompt_type_choices:
+        word_choices = ["test", "temp", "hello", "where"]
+        word = random.choice(word_choices)
+        if kind == "repeat":
+            prompt = f"""
+            please repeat the word '{word}' 10 times.
+            give no other output than the word at least ten times in a row,
+            in lowercase with spaces between each word and without quotes.
+            """
+        elif kind == "sentence":
+            prompt = f"""
+            please give a ten-word sentence that
+            uses the word {word} at least once.
+            give no other output than that simple sentence without quotes.
+            """
+        else:
+            raise ValueError(f"Unknown prompt type: {kind}")
+        prompts.append([{"role": "user", "content": prompt}])
+
+    return prompts
+
+
+@pytest.fixture
+def sampling_config():
+    return SamplingParams(temperature=0, max_tokens=256, ignore_eos=False)
+
+
+@pytest.fixture
+def model_name():
+    return "wemaster/deepseek_mtp_main_random_bf16"
+
+
+def test_mtp_correctness(
+    monkeypatch: pytest.MonkeyPatch,
+    test_prompts: list[list[dict[str, Any]]],
+    sampling_config: SamplingParams,
+    model_name: str,
+):
+    '''
+    Compare the outputs of a original LLM and a speculative LLM
+    should be the same when using mtp speculative decoding.
+    '''
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+
+        ref_llm = LLM(model=model_name, max_model_len=256)
+        ref_outputs = ref_llm.chat(test_prompts, sampling_config)
+        del ref_llm
+
+        spec_llm = LLM(model=model_name,
+                       trust_remote_code=True,
+                       speculative_config={
+                           "method": "deepseek_mtp",
+                           "num_speculative_tokens": 1,
+                       },
+                       max_model_len=256,
+                       enforce_eager=False)
+        spec_outputs = spec_llm.chat(test_prompts, sampling_config)
+        matches = 0
+        misses = 0
+        for ref_output, spec_output in zip(ref_outputs, spec_outputs):
+            if ref_output.outputs[0].text == spec_output.outputs[0].text:
+                matches += 1
+            else:
+                misses += 1
+                print(f"ref_output: {ref_output.outputs[0].text}")
+                print(f"spec_output: {spec_output.outputs[0].text}")
+
+        # Heuristic: expect at least 66% of the prompts to match exactly
+        # Upon failure, inspect the outputs to check for inaccuracy.
+        assert matches > int(0.66 * len(ref_outputs))
+        del spec_llm

vllm_ascend/attention/mla_v1.py

@@ -17,13 +17,26 @@
 from vllm_ascend.attention.attention_v1 import AscendAttentionState
 from vllm_ascend.ops.attention import vanilla_chunked_prefill_mla
 from vllm_ascend.utils import vllm_version_is
-from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
     from vllm.v1.worker.gpu_input_batch import InputBatch
 
 
+@dataclass
+class CommonAttentionMetadata:
+    """
+    Attention metadata attributes that can be shared by layers in different KV
+    cache groups and thus having different block table.
+    """
+
+    query_start_loc: torch.Tensor
+    """(batch_size + 1,), the start location of each request in query Tensor"""
+    seq_lens: torch.Tensor
+    """(batch_size,), the length of each request including both computed tokens
+    and newly scheduled tokens"""
+
+
 class AscendMLABackend(AttentionBackend):
 
     accept_output_buffer: bool = True
@@ -58,6 +71,7 @@ class AscendMLAPrefillMetadata:
     seq_lens: list[int]
     context_lens: torch.Tensor
     input_positions: torch.Tensor
+    query_start_loc: torch.Tensor
     block_table: torch.Tensor
     max_query_len: int
     max_seq_lens: int
@@ -91,6 +105,9 @@ class AscendMLAMetadata:
 
     num_actual_tokens: int  # Number of tokens excluding padding.
     slot_mapping: torch.Tensor
+    query_start_loc: torch.Tensor
+    seq_lens: torch.Tensor
+    block_tables: torch.Tensor
 
     # New for MLA (compared to FlashAttention)
     # For handling prefill decode split
@@ -131,7 +148,7 @@ class AscendMLAMetadataBuilder:
 
     # _attn_mask_builder = None
     def __init__(self,
-                 runner: "NPUModelRunner",
+                 runner,
                  metadata_cls: Optional[AscendMLAMetadata] = None):
         self.metadata_cls: Optional[AscendMLAMetadata] = metadata_cls \
             if metadata_cls is not None else AscendMLAMetadata  # type: ignore
@@ -231,6 +248,7 @@ def build(self,
               num_reqs: int,
               num_actual_tokens: int,
               max_query_len: int,
+              common_attn_metadata: CommonAttentionMetadata,
               common_prefix_len: Optional[int] = None,
               graph_pad_size: int = -1) -> AscendMLAMetadata:
         assert self._num_decodes + self._num_prefills == num_reqs
@@ -243,10 +261,8 @@ def build(self,
             block_table = (self.runner.input_batch.block_table.
                            get_device_tensor()[:num_reqs])
         else:
-            block_table = self.runner.input_batch.block_table[
-                0].get_device_tensor()
-            block_table[:num_reqs, :self.runner.max_num_blocks_per_req] = (
-                block_table[:num_reqs])
+            block_table = (self.runner.input_batch.block_table[0].
+                           get_device_tensor()[:num_reqs])
         slot_mapping = self.runner.slot_mapping_cpu[:num_actual_tokens].to(
             device, non_blocking=True)
         input_positions = self.runner.positions_cpu[:num_actual_tokens].to(
@@ -258,13 +274,17 @@ def build(self,
         seq_lens = seq_lens_cpu
         max_query_len = query_lens.max().item()
         max_seq_lens = seq_lens.max().item()
+        query_start_loc = None
 
         prefill_metadata = None
         if self._num_prefills > 0:
             reqs_start = self._num_decodes  # prefill_start
             tokens_start = self._num_decode_tokens
             max_query_len = query_lens[tokens_start:].max().item()
             max_seq_lens = seq_lens[tokens_start:].max().item()
+            query_start_loc = common_attn_metadata.query_start_loc
+            prefill_query_start_loc = query_start_loc[
+                reqs_start:] - query_start_loc[reqs_start]
 
             prefill_metadata = AscendMLAPrefillMetadata(
                 attn_mask=self.runner.attn_mask,
@@ -275,6 +295,7 @@ def build(self,
                 block_table=block_table[reqs_start:, ...],
                 max_query_len=max_query_len,
                 max_seq_lens=max_seq_lens,
+                query_start_loc=prefill_query_start_loc,
             )
 
         decode_metadata = None
@@ -331,6 +352,9 @@ def build(self,
             attn_state=self.runner.attn_state,
             prefill=prefill_metadata,
             decode=decode_metadata,
+            query_start_loc=query_start_loc,
+            block_tables=block_table,
+            seq_lens=seq_lens,
         )
 
 

vllm_ascend/ops/attention.py

@@ -222,6 +222,17 @@ def vanilla_chunked_prefill_mla(
         device="npu",
         dtype=value.dtype,
     )
+    num_query = torch.sum(q_mask).item()
+    num_add_query = num_query - query.size(0)
+    # mtp will come in
+    if num_add_query != 0:
+        add_query_size = query.size()
+        add_query_size = list(add_query_size)
+        add_query_size[0] = num_add_query
+        pad_tensor = torch.zeros(add_query_size,
+                                 dtype=query.dtype,
+                                 device=query.device)
+        query = torch.cat([query, pad_tensor], dim=0)
     pad_q[q_mask] = query
     pad_k[kv_c_mask] = key[kv_c_mask]
     pad_v[kv_c_mask] = value[kv_c_mask]
@@ -247,8 +258,8 @@ def vanilla_chunked_prefill_mla(
 
     attn_output = (attn_output[q_mask].view([-1, num_heads,
                                              v_head_dim]).to(output.dtype))
-    output = output.view_as(attn_output)
-    output.copy_(attn_output)
+    output = output.view([-1, num_heads, v_head_dim])
+    output.copy_(attn_output[:query.size(0) - num_add_query])
     return attn_output
 
 

vllm_ascend/patch/platform/patch_main/__init__.py

@@ -13,4 +13,4 @@
 # 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.
-#
+#