[BugFix][SpecDecode] Fix low MTP acceptance rate for SFA+DSA_CP (MTP>1) (vllm-project#10878)

### What this PR does / why we need it?
Fixes a degraded multi-token-prediction (MTP) / speculative-decoding
acceptance rate on Ascend when num_speculative_tokens > 1 when enable
dsa_cp for sfa models.

Root cause: Root cause: In the SFA v1 attention-metadata builder
(vllm_ascend/attention/sfa_v1.py), the actual_seq_lengths_query /
actual_seq_lengths_key buffers handed to lightning indexer shared the
same memory buffer across all draft steps. With MTP > 1, the
attention-metadata is overwritten and only the last step is kept, which
leads to wrong inputs. Meanwhile, rotary cos/sin were also fetched from
an out-dated cache while drafting.

Align Ascend attention builders' build_for_drafting with vllm's base
convention (common_attn_metadata, draft_index) so all spec-decode draft
paths go through one consistent entry point.

- sfa_v1.py: split build() into build()/build_for_drafting()/_build();
pre-allocate per-draft-step spec_actual_seq_lengths_query/key buffers so
each draft step owns independent metadata (no cross-step aliasing);
recompute rotary (use_cache=False) while drafting.
- llm_base_proposer.py: drop the else->build(0,...,model) branch; route
all draft steps through build_for_drafting(common_attn_metadata,
draft_index, **extra). Standard-attention drafters (eagle/eagle3/mtp on
Qwen/Llama) now hit the base build_for_drafting -> build(0, m,
fast_build=True), fixing the 'int' object has no attribute 'num_reqs'
crash from the swapped-positional elif. Rename draft_step ->
draft_index.
- dsa_v1.py / dsa_cp.py: swap build_for_drafting arg order to
(common_attn_metadata, draft_index, ...) and rename draft_step ->
draft_index for consistency.

- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@967c5c3

Co-authored-by: wu-yushan <wuyushan1@huawei.com>
Co-authored-by: YzTongNiar <1667927948@qq.com>
Co-authored-by: Orla-see <lidan276@huawei.com>

Author: 4 people

vllm_ascend/attention/context_parallel/dsa_cp.py

@@ -352,8 +352,8 @@ def build(
 
     def build_for_drafting(
         self,
-        draft_step: int,
         common_attn_metadata: AscendCommonAttentionMetadata,
+        draft_index: int,
         fast_build: bool = False,
         **kwargs,
     ) -> AscendDSAMetadata:
@@ -381,13 +381,13 @@ def build_for_drafting(
         slot_mapping = common_attn_metadata.slot_mapping[:num_input_tokens]
 
         assert self.spec_slot_mapping is not None
-        self.spec_slot_mapping[draft_step - 1][:num_input_tokens] = DeviceOperator.format_dsa_slot_mapping(
+        self.spec_slot_mapping[draft_index - 1][:num_input_tokens] = DeviceOperator.format_dsa_slot_mapping(
             slot_mapping, self.block_size
         )
 
         self.block_table = common_attn_metadata.block_table_tensor[:num_reqs]
         req_metadata = self.build_req_metadata_for_drafting(
-            draft_step=draft_step,
+            draft_index=draft_index,
             common_attn_metadata=common_attn_metadata,
             input_positions=input_positions,
             num_input_tokens=num_input_tokens,
@@ -413,7 +413,7 @@ def build_for_drafting(
 
     def build_req_metadata_for_drafting(
         self,
-        draft_step: int,
+        draft_index: int,
         common_attn_metadata: AscendCommonAttentionMetadata,
         input_positions: torch.Tensor,
         num_input_tokens: int,
@@ -442,8 +442,8 @@ def build_req_metadata_for_drafting(
             query_start_loc=query_start_loc,
             seq_lens=self.seq_lens[:num_reqs],
             use_cache=False,
-            local_query_start_loc=self.spec_local_query_start_loc[draft_step - 1],
-            local_seq_lens=self.spec_local_seq_lens[draft_step - 1],
+            local_query_start_loc=self.spec_local_query_start_loc[draft_index - 1],
+            local_seq_lens=self.spec_local_seq_lens[draft_index - 1],
         )
         local_query_start_loc = local_query_start_loc.clone()
         local_seq_lens = local_seq_lens.clone()
@@ -463,7 +463,7 @@ def build_req_metadata_for_drafting(
         start_pos = self.seq_lens[:num_reqs] - seq_lens_q
 
         assert self.spec_slot_mapping is not None
-        slot_mapping = self.spec_slot_mapping[draft_step - 1][: self.num_actual_tokens]
+        slot_mapping = self.spec_slot_mapping[draft_index - 1][: self.num_actual_tokens]
 
         num_heads = self.model_config.hf_config.num_attention_heads
         metadata_op = DeviceOperator.get_dsa_sparse_attn_metadata_op()

vllm_ascend/attention/dsa_v1.py

@@ -1133,8 +1133,8 @@ def _get_compressed_decode_token_start(decode_input_positions, compress_ratio):
 
     def build_for_drafting(
         self,
-        draft_step: int,
         common_attn_metadata: AscendCommonAttentionMetadata,
+        draft_index: int,
         fast_build: bool = False,
         **kwargs,
     ) -> AscendDSADecodeMetadata:
@@ -1147,19 +1147,19 @@ def build_for_drafting(
         if num_prefills:
             cos, sin = get_cos_and_sin_dsa(input_positions)
         else:
-            # disable use_cache, otherwise, draft_step>0 will override draft_step=0
+            # disable use_cache, otherwise, draft_index>0 will override draft_index=0
             # take care of this, if full graph is needed then rope cache is inevitable
             cos, sin = get_cos_and_sin_dsa(input_positions, use_cache=False)
 
         slot_mapping = common_attn_metadata.slot_mapping[:num_input_tokens]
-        self.spec_slot_mapping[draft_step - 1][:num_input_tokens] = DeviceOperator.format_dsa_slot_mapping(  # type: ignore[index]
+        self.spec_slot_mapping[draft_index - 1][:num_input_tokens] = DeviceOperator.format_dsa_slot_mapping(  # type: ignore[index]
             slot_mapping, self.block_size
         )
 
         prefill_metadata = None
         if num_prefills > 0:
             prefill_metadata = self.build_prefill_metadata_for_drafting(
-                draft_step=draft_step,
+                draft_index=draft_index,
                 common_attn_metadata=common_attn_metadata,
                 reqs_start=num_decodes,
                 tokens_start=num_decode_tokens,
@@ -1169,7 +1169,7 @@ def build_for_drafting(
         decode_metadata = None
         if num_decodes > 0:
             decode_metadata = self.build_decode_metadata_for_drafting(
-                draft_step=draft_step,
+                draft_index=draft_index,
                 common_attn_metadata=common_attn_metadata,
                 num_decodes=num_decodes,
                 num_decode_tokens=num_decode_tokens,
@@ -1198,7 +1198,7 @@ def build_for_drafting(
 
     def build_prefill_metadata_for_drafting(
         self,
-        draft_step: int,
+        draft_index: int,
         common_attn_metadata: AscendCommonAttentionMetadata,
         **kwargs,
     ) -> AscendDSAPrefillMetadata:
@@ -1218,7 +1218,7 @@ def build_prefill_metadata_for_drafting(
         prefill_input_positions = input_positions[tokens_start:]
         cos, sin = get_cos_and_sin_dsa(prefill_input_positions)
 
-        prefill_slot_mapping = self.spec_slot_mapping[draft_step - 1][tokens_start:num_prefill_tokens]  # type: ignore[index]
+        prefill_slot_mapping = self.spec_slot_mapping[draft_index - 1][tokens_start:num_prefill_tokens]  # type: ignore[index]
         block_table = common_attn_metadata.block_table_tensor[: common_attn_metadata.num_reqs]
 
         metadata_op = DeviceOperator.get_dsa_sparse_attn_metadata_op()
@@ -1270,7 +1270,7 @@ def build_prefill_metadata_for_drafting(
 
     def build_decode_metadata_for_drafting(
         self,
-        draft_step: int,
+        draft_index: int,
         common_attn_metadata: AscendCommonAttentionMetadata,
         **kwargs,
     ) -> AscendDSADecodeMetadata:
@@ -1295,11 +1295,11 @@ def build_decode_metadata_for_drafting(
         max_seqlen_kv = torch.max(_seq_lens_cpu[:num_decodes]).item()
 
         input_positions = common_attn_metadata.positions[:num_decode_tokens_typed].long()
-        # disable use_cache, otherwise, draft_step>0 will override draft_step=0
+        # disable use_cache, otherwise, draft_index>0 will override draft_index=0
         # take care of this, if full graph is needed then rope cache is inevitable
         cos, sin = get_cos_and_sin_dsa(input_positions, use_cache=False)
 
-        slot_mapping = self.spec_slot_mapping[draft_step - 1][:num_decode_tokens_typed]  # type: ignore[index]
+        slot_mapping = self.spec_slot_mapping[draft_index - 1][:num_decode_tokens_typed]  # type: ignore[index]
         block_table = common_attn_metadata.block_table_tensor
 
         metadata_op = DeviceOperator.get_dsa_sparse_attn_metadata_op()

vllm_ascend/attention/sfa_v1.py

@@ -200,6 +200,11 @@ def __init__(
 
         self.speculative_config = vllm_config.speculative_config
         self.decode_threshold = 1
+        max_num_reqs = vllm_config.scheduler_config.max_num_seqs
+        self.actual_seq_lengths_query = torch.zeros(max_num_reqs + 1, dtype=torch.int32, device=device)
+        self.actual_seq_lengths_key = torch.empty_like(self.actual_seq_lengths_query)
+        self.spec_actual_seq_lengths_query: list[torch.Tensor] | None = None
+        self.spec_actual_seq_lengths_key: list[torch.Tensor] | None = None
         if self.speculative_config:
             spec_token_num = self.speculative_config.num_speculative_tokens
             self.decode_threshold += spec_token_num
@@ -208,15 +213,20 @@ def __init__(
                 npu_fused_infer_attention_score TND layout's limit of 16, \
                 got {self.decode_threshold}"
             )
+            self.spec_actual_seq_lengths_query = [
+                torch.zeros(max_num_reqs * (spec_token_num + 1) + 1, dtype=torch.int32, device=device)
+                for _ in range(spec_token_num)
+            ]
+            self.spec_actual_seq_lengths_key = [
+                torch.zeros(max_num_reqs * (spec_token_num + 1) + 1, dtype=torch.int32, device=device)
+                for _ in range(spec_token_num)
+            ]
+
         self.reorder_batch_threshold = self.decode_threshold
         self.attn_mask_builder = AttentionMaskBuilder(self.device)
         self.rope_dim = self.model_config.hf_text_config.qk_rope_head_dim
         self.enable_dsa_cp = enable_dsa_cp()
 
-        max_num_reqs = vllm_config.scheduler_config.max_num_seqs
-        self.actual_seq_lengths_query = torch.zeros(max_num_reqs + 1, dtype=torch.int32, device=device)
-        self.actual_seq_lengths_key = torch.empty_like(self.actual_seq_lengths_query)
-
     @staticmethod
     def determine_chunked_prefill_workspace_size(vllm_config: VllmConfig) -> int:
         return ascend_chunked_prefill_workspace_size(vllm_config)
@@ -240,6 +250,22 @@ def build(
         common_prefix_len: int,
         common_attn_metadata: AscendCommonAttentionMetadata,
         fast_build: bool = False,
+    ) -> AscendSFAMetadata:
+        # common_prefix_len / fast_build are unused; kept for API compatibility.
+        return self._build(common_attn_metadata, draft_index=None)
+
+    def build_for_drafting(
+        self,
+        common_attn_metadata: AscendCommonAttentionMetadata,
+        draft_index: int,
+        **kwargs,
+    ) -> AscendSFAMetadata:
+        return self._build(common_attn_metadata, draft_index=draft_index)
+
+    def _build(
+        self,
+        common_attn_metadata: AscendCommonAttentionMetadata,
+        draft_index: int | None = None,
     ) -> AscendSFAMetadata:
         num_reqs = common_attn_metadata.num_reqs
         num_actual_tokens = common_attn_metadata.num_actual_tokens
@@ -265,7 +291,7 @@ def build(
         else:
             seq_lens_cpu = common_attn_metadata.seq_lens[:num_reqs].to("cpu")
 
-        cos, sin = get_cos_and_sin_mla(input_positions, True)
+        cos, sin = get_cos_and_sin_mla(input_positions, use_cache=(draft_index is None))
 
         dsa_cp_context = None
         if self.enable_dsa_cp:
@@ -307,8 +333,16 @@ def build(
                     got {slot_mapping.shape[0]} and {num_tokens_pad}"
             )
 
-            actual_seq_lengths_query = self.actual_seq_lengths_query
-            actual_seq_lengths_key = self.actual_seq_lengths_key
+            if draft_index is not None:
+                assert self.spec_actual_seq_lengths_query is not None
+                assert self.spec_actual_seq_lengths_key is not None
+                # Per-draft-step buffers: independent, graph-stable storage so
+                # later draft steps don't clobber earlier ones' metadata.
+                actual_seq_lengths_query = self.spec_actual_seq_lengths_query[draft_index - 1]
+                actual_seq_lengths_key = self.spec_actual_seq_lengths_key[draft_index - 1]
+            else:
+                actual_seq_lengths_query = self.actual_seq_lengths_query
+                actual_seq_lengths_key = self.actual_seq_lengths_key
 
             num_segs = cum_query_lens.shape[0]
             last_token = 0