[Feature] Support dsv4 mtp graph (vllm-project#11062)

### What this PR does / why we need it?
This PR aims to supoort dsv4 mtp graph. The main change is adding proper
arguments in llm_base_proposer dummy_run. As for rope, we additionaly
create buffers for different draft steps, solving low acceptance rate
issue. Furthermore, _pad_query_start_loc_for_fia is modified so that
target and draft model have their own variables and avoid race when
using copy_to_gpu.

### Does this PR introduce _any_ user-facing change?
N/A

### How was this patch tested?
by ci

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

---------

Signed-off-by: Zetong Li <slippersss@126.com>

Author: slippersss

tests/e2e/pull_request/four_card/test_deepseek_v4.py

@@ -55,7 +55,7 @@ def test_deepseek_v4_w4a8_tp4_basic_greedy():
         compilation_config={
             "cudagraph_mode": "FULL_DECODE_ONLY",
         },
-        speculative_config={"num_speculative_tokens": 1, "method": "mtp", "enforce_eager": True},
+        speculative_config={"num_speculative_tokens": 1, "method": "mtp"},
     ) as vllm_model:
         outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
         expected_token_ids = [

tests/ut/spec_decode/a2/test_eagle_proposer.py

@@ -1160,6 +1160,9 @@ def test_propose(self, mock_get_model, graphmode, model_type, flag_prefill_decod
         original_method = self.proposer.attn_update_stack_num_spec_norm
         mock_bd = MagicMock()
         mock_bd.num_tokens = 16
+        self.proposer.query_start_loc = MagicMock()
+        self.proposer.query_start_loc.gpu = torch.tensor([0, 4, 8, 12, 16], device=torch.device("cpu"), dtype=torch.int32)
+        self.proposer.query_start_loc.cpu = torch.tensor([0, 4, 8, 12, 16], device=torch.device("cpu"), dtype=torch.int32)
         self.runner.cudagraph_dispatcher.dispatch.return_value = (CUDAGraphMode.FULL, mock_bd)
         self.runner._pad_query_start_loc_for_fia.return_value = 4
         self.runner.query_start_loc.gpu = torch.tensor([0, 4, 8, 12, 16], device=torch.device("cpu"), dtype=torch.int32)
@@ -1583,7 +1586,7 @@ def check_mock(self):
         assert hasattr(RunnerCls, "_pad_query_start_loc_for_fia")
         sig = inspect.signature(RunnerCls._pad_query_start_loc_for_fia)
         sig_name = self.get_param_names(sig)
-        assert sig_name == ['self', 'num_tokens_padded', 'num_reqs_padded', 'num_reqs', 'cudagraph_runtime_mode', 'batch_desc_num_reqs']
+        assert sig_name == ['self', 'query_start_loc', 'num_tokens_padded', 'num_reqs_padded', 'num_reqs', 'cudagraph_runtime_mode', 'batch_desc_num_reqs']
 
 
         import vllm_ascend.spec_decode.llm_base_proposer

vllm_ascend/_310p/model_runner_310p.py

@@ -197,6 +197,7 @@ def _build_attention_metadata(self, *args: Any, **kwargs: Any):
 
     def _pad_query_start_loc_for_fia(
         self,
+        query_start_loc: torch.Tensor,
         num_tokens_padded: int,
         num_reqs_padded: int,
         num_reqs: int,
@@ -213,19 +214,19 @@ def _pad_query_start_loc_for_fia(
             # Uniform-batch case: num_reqs must be no greater than num_reqs_padded
             assert num_reqs <= num_reqs_padded
 
-            last_loc = self.query_start_loc.np[num_reqs]
-            self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1] = (
+            last_loc = query_start_loc.np[num_reqs]
+            query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1] = (
                 self.arange_np[1 : num_reqs_padded + 1 - num_reqs] * uniform_decode_query_len + last_loc
             )
         else:
             # Mixed-batch case: num_reqs must equal num_reqs_padded
             assert num_reqs == num_reqs_padded
 
             # Insert a dummy request instead of setting query_start_loc[num_reqs] = num_tokens_padded directly
-            self.query_start_loc.np[num_reqs_padded + 1] = num_tokens_padded
+            query_start_loc.np[num_reqs_padded + 1] = num_tokens_padded
             num_reqs_padded = num_reqs_padded + 1
 
-        self.query_start_loc.copy_to_gpu()
+        query_start_loc.copy_to_gpu()
         return num_reqs_padded
 
     def _build_attn_state(self, num_reqs, num_scheduled_tokens, num_valid_tokens):

vllm_ascend/attention/dsa_v1.py

@@ -397,6 +397,9 @@ def __init__(
                 torch.zeros(self.slot_mapping_shape, dtype=torch.int32, device=self.device)
                 for _ in range(spec_token_num)
             ]
+            self.spec_sas_metadata = [
+                torch.zeros(1024, dtype=torch.int32, device=self.device) for _ in range(spec_token_num)
+            ]
             self.decode_threshold += spec_token_num
             assert self.decode_threshold <= 16, (
                 f"decode_threshold exceeded \
@@ -1111,7 +1114,7 @@ def build_for_drafting(
         else:
             # 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)
+            cos, sin = get_cos_and_sin_dsa(input_positions, use_cache=True, draft_index=draft_index)
 
         slot_mapping = common_attn_metadata.slot_mapping[:num_input_tokens]
         self.spec_slot_mapping[draft_index - 1][:num_input_tokens] = DeviceOperator.format_dsa_slot_mapping(  # type: ignore[index]
@@ -1258,7 +1261,7 @@ def build_decode_metadata_for_drafting(
         input_positions = common_attn_metadata.positions[:num_decode_tokens_typed].long()
         # 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)
+        cos, sin = get_cos_and_sin_dsa(input_positions, use_cache=True, draft_index=draft_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
@@ -1289,6 +1292,8 @@ def build_decode_metadata_for_drafting(
             has_ori_kv=True,
             has_cmp_kv=False,
         )
+        self.spec_sas_metadata[draft_index - 1][:1024].copy_(decode_sas_metadata[:1024])
+        decode_sas_metadata = self.spec_sas_metadata[draft_index - 1]
 
         decode_metadata = AscendDSADecodeMetadata(
             input_positions=None,
@@ -1452,6 +1457,19 @@ def __init__(
             False,
         )
 
+    @staticmethod
+    def update_graph_params(
+        update_stream,
+        forward_context,
+        num_tokens,
+        vllm_config=None,
+        speculative_config=None,
+        num_dcp_pcp_tokens=None,
+        draft_attn_metadatas=None,
+    ):
+        # dsa does not need to update graph params
+        pass
+
     def _get_indexcache_topk_indices(self, num_tokens: int, offset: int = 0) -> torch.Tensor:
         if self.topk_indices_buffer is None:
             raise RuntimeError("IndexCache requires topk_indices_buffer when skip_topk is enabled.")

vllm_ascend/ops/rope_dsv4.py

@@ -12,6 +12,7 @@ class RopeGlobalState:
     def __init__(self):
         self.full_rope_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
         self.runtime_buffer: dict[str, dict[str, tuple[torch.Tensor, torch.Tensor]]] = {}
+        self.spec_runtime_buffer: dict[str, dict[str, tuple[torch.Tensor, torch.Tensor]]] = {}
         self.layer_info: dict[str, tuple[str, list[str]]] = {}
         self.registry_summary: dict[str, set] = {}
 
@@ -58,7 +59,11 @@ def __getitem__(self, index):
             return layer_result
 
 
-def get_cos_and_sin_dsa(positions: torch.Tensor | dict[str, torch.Tensor], use_cache: bool = False):
+def get_cos_and_sin_dsa(
+    positions: torch.Tensor | dict[str, torch.Tensor],
+    use_cache: bool = False,
+    draft_index: int | None = None,
+):
     if isinstance(positions, torch.Tensor):
         pos_map = {"default": positions}
     else:
@@ -81,18 +86,30 @@ def get_cos_and_sin_dsa(positions: torch.Tensor | dict[str, torch.Tensor], use_c
             curr_sin = full_rope_sin[pos_tensor]
 
             if use_cache:
-                group_buffers = _ROPE_STATE.runtime_buffer.get(config_key, {}).get(group_name)
+                group_buffers = (
+                    _ROPE_STATE.runtime_buffer.get(config_key, {}).get(group_name)
+                    if draft_index is None
+                    else _ROPE_STATE.spec_runtime_buffer.get(config_key, {}).get(group_name)
+                )
 
                 if group_buffers is None:
                     continue
 
                 buf_cos, buf_sin = group_buffers
                 num_tokens = pos_tensor.size(0)
 
-                buf_cos[:num_tokens].copy_(curr_cos)
-                buf_sin[:num_tokens].copy_(curr_sin)
+                if draft_index is None:
+                    buf_cos[:num_tokens].copy_(curr_cos)
+                    buf_sin[:num_tokens].copy_(curr_sin)
 
-                batch_result[config_key][group_name] = (buf_cos[:num_tokens], buf_sin[:num_tokens])
+                    batch_result[config_key][group_name] = (buf_cos[:num_tokens], buf_sin[:num_tokens])
+                else:
+                    buf_cos[draft_index - 1][:num_tokens].copy_(curr_cos)
+                    buf_sin[draft_index - 1][:num_tokens].copy_(curr_sin)
+                    batch_result[config_key][group_name] = (
+                        buf_cos[draft_index - 1][:num_tokens],
+                        buf_sin[draft_index - 1][:num_tokens],
+                    )
             else:
                 batch_result[config_key][group_name] = (curr_cos, curr_sin)
 
@@ -171,8 +188,17 @@ def __init__(
 
             _ROPE_STATE.full_rope_cache[config_key] = (cos.unsqueeze(1).unsqueeze(1), sin.unsqueeze(1).unsqueeze(1))
 
+        use_eagle = (
+            vllm_config is not None
+            and vllm_config.speculative_config is not None
+            and vllm_config.speculative_config.use_eagle()
+        )
+        num_speculative_tokens = vllm_config.speculative_config.num_speculative_tokens if use_eagle else None
+
         if config_key not in _ROPE_STATE.runtime_buffer:
             _ROPE_STATE.runtime_buffer[config_key] = {}
+            if num_speculative_tokens is not None:
+                _ROPE_STATE.spec_runtime_buffer[config_key] = {}
 
         target_device = current_platform.device_type
         max_batch_size = vllm_config.scheduler_config.max_num_batched_tokens
@@ -181,6 +207,16 @@ def __init__(
                 buf_cos = torch.ones(max_batch_size, 1, 1, rotary_dim, dtype=torch.float32, device=target_device)
                 buf_sin = torch.zeros(max_batch_size, 1, 1, rotary_dim, dtype=torch.float32, device=target_device)
                 _ROPE_STATE.runtime_buffer[config_key][grp] = (buf_cos, buf_sin)
+                if num_speculative_tokens is not None:
+                    buf_cos = [
+                        torch.ones(max_batch_size, 1, 1, rotary_dim, dtype=torch.float32, device=target_device)
+                        for _ in range(num_speculative_tokens)
+                    ]
+                    buf_sin = [
+                        torch.zeros(max_batch_size, 1, 1, rotary_dim, dtype=torch.float32, device=target_device)
+                        for _ in range(num_speculative_tokens)
+                    ]
+                    _ROPE_STATE.spec_runtime_buffer[config_key][grp] = (buf_cos, buf_sin)
 
     @staticmethod
     def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow):

vllm_ascend/spec_decode/llm_base_proposer.py

@@ -551,6 +551,7 @@ def dummy_run(
                 query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
                 query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
                 seq_lens_cpu=self.runner.optimistic_seq_lens_cpu,
+                _seq_lens_cpu=self.runner.optimistic_seq_lens_cpu,
                 seq_lens_cpu_upper_bound=self.runner.optimistic_seq_lens_cpu,
                 seq_lens=self.runner.seq_lens[:num_reqs],
                 num_reqs=num_reqs,
@@ -559,11 +560,14 @@ def dummy_run(
                 max_query_len=self.num_speculative_tokens + 1,
                 num_computed_tokens_cpu=num_computed_tokens_cpu,
                 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],
+                block_table_tensor=self.runner.input_batch.block_table[self.kv_cache_gid].get_device_tensor()[
+                    :num_reqs
+                ],
                 # This is used to hold a position.
-                slot_mapping=self.runner.input_batch.block_table[0].slot_mapping.gpu,
-                slot_mapping_cpu=self.runner.input_batch.block_table[0].slot_mapping.cpu,
+                slot_mapping=self.runner.input_batch.block_table[self.kv_cache_gid].slot_mapping.gpu,
+                slot_mapping_cpu=self.runner.input_batch.block_table[self.kv_cache_gid].slot_mapping.cpu,
                 positions=self.runner.positions,
+                positions_cpu=self.runner._dsa_positions_cpu_buf if self.use_compress else None,
                 attn_state=self.runner.attn_state,
                 decode_token_per_req=self.runner.decode_token_per_req,
                 is_prefilling=torch.zeros(num_reqs, dtype=torch.bool),
@@ -575,20 +579,46 @@ def dummy_run(
 
             assert len(self.draft_attn_groups) > 0
             builder = self.draft_attn_groups[0].get_metadata_builder()
+            kv_cache_spec = self.draft_attn_groups[0].kv_cache_spec
             # update the tensor's address for each step.
             for draft_index in range(self.num_speculative_tokens):
                 common_attn_metadata = self.shallow_copy_metadata(common_attn_metadata)
+                extra_attn_metadata_args: dict = {}
+                if self.use_compress:
+                    extra_attn_metadata_args = dict(
+                        prefill_ratio_to_sas_metadata=dict(),
+                        decode_ratio_to_sas_metadata=dict(),
+                        common_ratio_to_sas_metadata=dict(),
+                        block_size=kv_cache_spec.block_size,
+                    )
                 # Set the real slot_mapping.
+                slot_mapping_lens = common_attn_metadata.slot_mapping.shape[0]
+                self.slot_mapping_group[draft_index][:slot_mapping_lens].copy_(common_attn_metadata.slot_mapping)
+                self.slot_mapping_group[draft_index][slot_mapping_lens:].fill_(PADDING_SLOT_ID)
                 common_attn_metadata.slot_mapping = self.slot_mapping_group[draft_index]
+                self.seq_lens_group[draft_index][:num_reqs].copy_(common_attn_metadata.seq_lens)
+                self.seq_lens_group[draft_index][num_reqs:].fill_(0)
                 common_attn_metadata.seq_lens = self.seq_lens_group[draft_index][:num_reqs]
+                self.query_start_loc_group[draft_index][: num_reqs + 1].copy_(common_attn_metadata.query_start_loc)
+                self.query_start_loc_group[draft_index][num_reqs + 1 :].fill_(0)
                 common_attn_metadata.query_start_loc = self.query_start_loc_group[draft_index][: num_reqs + 1]
                 if self.pcp_size * self.dcp_size > 1 and draft_index > 0:
                     assert self.block_table_tensor_clone is not None, "block_table_tensor_clone is not init"
                     common_attn_metadata.block_table_tensor = self.block_table_tensor_clone[:num_reqs]
-                attn_metadata_eagle = builder.build_for_graph_capture(
-                    common_attn_metadata,
-                    AscendAttentionState.SpecDecoding if self.method == "mtp" else AscendAttentionState.ChunkedPrefill,
-                )
+                if not self.use_compress or draft_index == 0:
+                    attn_metadata_eagle = builder.build_for_graph_capture(
+                        common_attn_metadata,
+                        AscendAttentionState.SpecDecoding
+                        if self.method == "mtp"
+                        else AscendAttentionState.ChunkedPrefill,
+                        **extra_attn_metadata_args,
+                    )
+                else:
+                    attn_metadata_eagle = builder.build_for_drafting(
+                        common_attn_metadata,
+                        draft_index,
+                        **extra_attn_metadata_args,
+                    )
                 per_layer_attn_metadata = dict()
                 for layer_name in self.attn_layer_names:
                     per_layer_attn_metadata[layer_name] = attn_metadata_eagle
@@ -744,16 +774,20 @@ def _propose(
             # is run in eager mode currently, which means `_pad_query_start_loc_for_fia` is not called,
             # while draft model is run in graph model, which means we should pad the `query_start_loc`.
             # Need to be fixed in the future.
+            num_reqs = common_attn_metadata.query_start_loc.shape[0]
+            self.query_start_loc.gpu[:num_reqs].copy_(common_attn_metadata.query_start_loc)
+            self.query_start_loc.cpu[:num_reqs].copy_(common_attn_metadata.query_start_loc_cpu)
             num_reqs_padded = self.runner._pad_query_start_loc_for_fia(
+                self.query_start_loc,
                 num_input_tokens,
                 batch_descriptor.num_reqs if batch_descriptor.num_reqs is not None else common_attn_metadata.num_reqs,
                 common_attn_metadata.num_reqs,
                 aclgraph_runtime_mode,
                 batch_descriptor.num_reqs,
             )
             common_attn_metadata.num_reqs = num_reqs_padded
-            common_attn_metadata.query_start_loc = self.runner.query_start_loc.gpu[: num_reqs_padded + 1]
-            common_attn_metadata.query_start_loc_cpu = self.runner.query_start_loc.cpu[: num_reqs_padded + 1]
+            common_attn_metadata.query_start_loc = self.query_start_loc.gpu[: num_reqs_padded + 1]
+            common_attn_metadata.query_start_loc_cpu = self.query_start_loc.cpu[: num_reqs_padded + 1]
             slicing_length = (
                 num_reqs_padded * self.decode_threshold if self.pcp_size * self.dcp_size > 1 else num_reqs_padded
             )

vllm_ascend/worker/model_runner_v1.py

@@ -734,6 +734,7 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None
 
     def _pad_query_start_loc_for_fia(
         self,
+        query_start_loc: torch.Tensor,
         num_tokens_padded: int,
         num_reqs_padded: int,
         num_reqs: int,
@@ -762,21 +763,21 @@ def _pad_query_start_loc_for_fia(
             # Uniform-batch case: num_reqs must be no greater than num_reqs_padded
             assert num_reqs <= num_reqs_padded
 
-            last_loc = self.query_start_loc.np[num_reqs]
-            self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1] = (
+            last_loc = query_start_loc.np[num_reqs]
+            query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1] = (
                 self.arange_np[1 : num_reqs_padded + 1 - num_reqs] * self.uniform_decode_query_len + last_loc
             )
         else:
             # Mixed-batch case: num_reqs must equal num_reqs_padded
             assert num_reqs == num_reqs_padded
 
             # Do not insert if the last value already equals the num_tokens
-            if self.query_start_loc.np[num_reqs_padded] < num_tokens_padded:
+            if query_start_loc.np[num_reqs_padded] < num_tokens_padded:
                 # Insert a dummy request instead of change the last value directly
-                self.query_start_loc.np[num_reqs_padded + 1] = num_tokens_padded
+                query_start_loc.np[num_reqs_padded + 1] = num_tokens_padded
                 num_reqs_padded = num_reqs_padded + 1
 
-        self.query_start_loc.copy_to_gpu()
+        query_start_loc.copy_to_gpu()
 
         return num_reqs_padded
 
@@ -2208,7 +2209,12 @@ def execute_model(
                     # Another possible condition is num_tokens_padded != num_tokens_unpadded
                     # but this scope is way too big and the consequences are unpredictable
                     num_reqs_padded = self._pad_query_start_loc_for_fia(
-                        num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_mode, batch_desc.num_reqs
+                        self.query_start_loc,
+                        num_tokens_padded,
+                        num_reqs_padded,
+                        num_reqs,
+                        cudagraph_mode,
+                        batch_desc.num_reqs,
                     )
 
                 (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata(
@@ -3486,7 +3492,12 @@ def _dummy_run(
 
             if not profile_cpp:
                 num_reqs_padded = self._pad_query_start_loc_for_fia(
-                    num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_runtime_mode, batch_desc.num_reqs
+                    self.query_start_loc,
+                    num_tokens_padded,
+                    num_reqs_padded,
+                    num_reqs,
+                    cudagraph_runtime_mode,
+                    batch_desc.num_reqs,
                 )
 
             # Dummy graph runs do not go through _prepare_inputs(), but GDN/Mamba