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

### 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.

Changes:

- vllm_ascend/attention/sfa_v1.py: pre-allocate per-draft-step buffers
spec_actual_seq_lengths_query/key (one tensor per step, sized
max_num_reqs * (num_speculative_tokens + 1) + 1) so each step owns
independent, with no step overwrites another's metadata. Split the
builder into build() (non-draft/verify path),
build_for_drafting(draft_step, ...), and a shared _build(...,
draft_step) that selects the correct per-step buffer; pass
use_cache=False to get_cos_and_sin_mla while drafting so the rotary
tables are recomputed for each step.
- vllm_ascend/spec_decode/llm_base_proposer.py: when a builder exposes
build_for_drafting, route each draft step's attention-metadata build
through it instead of calling build(0, ...) for every step, so the
per-step buffers are actually used and cross-step aliasing is avoided.

Effect: correct per-step attention metadata and rotary embeddings →
restored MTP/spec-decode acceptance rate for num_speculative_tokens > 1.

### Does this PR introduce _any_ user-facing change?
No

### How was this patch tested?


- vLLM version: v0.22.1
- 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>
Co-authored-by: LookAround0301 <lixushi@huawei.com>

Author: 5 people

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_step=None)
+
+    def build_for_drafting(
+        self,
+        draft_step: int,
+        common_attn_metadata: AscendCommonAttentionMetadata,
+        **kwargs,
+    ) -> AscendSFAMetadata:
+        return self._build(common_attn_metadata, draft_step=draft_step)
+
+    def _build(
+        self,
+        common_attn_metadata: AscendCommonAttentionMetadata,
+        draft_step: 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_step 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_step 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_step - 1]
+                actual_seq_lengths_key = self.spec_actual_seq_lengths_key[draft_step - 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

vllm_ascend/spec_decode/llm_base_proposer.py

@@ -1652,6 +1652,13 @@ def attn_update_stack_num_spec_norm(
                 common_attn_metadata,
                 **extra_attn_metadata_args,
             )
+        elif hasattr(attn_metadata_builder, "build_for_drafting"):
+            # e.g. SFA (dsa_cp): route draft steps through a draft-aware build so
+            # per-draft-step buffers are used and cross-step aliasing is avoided.
+            attn_metadata = attn_metadata_builder.build_for_drafting(
+                draft_step,
+                common_attn_metadata,
+            )
         else:
             attn_metadata = attn_metadata_builder.build(
                 0,