diff --git a/tests/ut/worker/test_mtp_pcp_speculator_v2.py b/tests/ut/worker/test_mtp_pcp_speculator_v2.py index 23e2a9b0cb59..1f51f64b16bb 100644 --- a/tests/ut/worker/test_mtp_pcp_speculator_v2.py +++ b/tests/ut/worker/test_mtp_pcp_speculator_v2.py @@ -294,40 +294,46 @@ def test_prefill_rebuilds_replicated_pcp_metadata_before_filtering() -> None: assert parent_args[3] is global_slot_mappings -@pytest.mark.parametrize("replicated_pcp", [False, True]) -def test_graph_prefill_builds_draft_metadata(replicated_pcp: bool) -> None: +@pytest.mark.parametrize( + ("attn_architecture", "replicated_pcp", "rebuild_metadata"), + [ + ("DSA", True, True), + ("DSA", False, False), + ("SFA", True, True), + ("MLA", True, False), + ], +) +def test_graph_prefill_builds_draft_metadata( + attn_architecture: str, replicated_pcp: bool, rebuild_metadata: bool +) -> None: speculator = object.__new__(AscendMTPSpeculator) speculator.replicated_pcp = replicated_pcp + speculator.attn_architecture = attn_architecture + speculator.input_batch = _make_padded_input_batch() + speculator.input_batch.is_dummy = False + speculator.block_tables = MagicMock() + speculator.kv_cache_config = object() speculator.draft_attn_layer_names = {"draft.layer"} - local_draft_metadata = object() + local_draft_metadata = SimpleNamespace(decode=SimpleNamespace(actual_seq_lengths_q=[4, 8])) global_draft_metadata = object() speculator.model_state = SimpleNamespace( - attn_metadata={ - "draft.layer": local_draft_metadata, - "target.layer": object(), - } + attn_metadata={"draft.layer": local_draft_metadata, "target.layer": object()}, ) - prepared_draft_metadata = global_draft_metadata if replicated_pcp else local_draft_metadata - speculator._prepare_replicated_prefill_attn = MagicMock( - return_value=( - {"draft.layer": prepared_draft_metadata}, - None, - ) + speculator._build_draft_attn_metadata = MagicMock( + return_value={"draft.layer": global_draft_metadata}, ) - actual = speculator.build_draft_attn_metadatas( - num_reqs_padded=4, - num_tokens_padded=8, - is_draft_model_prefill=True, - ) + with patch.object(speculator_module, "build_slot_mappings_by_layer", return_value={}): + actual = speculator.build_draft_attn_metadatas( + num_reqs_padded=2, + num_tokens_padded=8, + is_draft_model_prefill=True, + ) - assert actual == [{"draft.layer": prepared_draft_metadata}] - speculator._prepare_replicated_prefill_attn.assert_called_once_with( - {"draft.layer": local_draft_metadata}, - None, - 4, - 8, - ) + expected_metadata = global_draft_metadata if rebuild_metadata else local_draft_metadata + assert actual == [{"draft.layer": expected_metadata}] + assert speculator._build_draft_attn_metadata.call_count == int(rebuild_metadata) + assert local_draft_metadata.decode.actual_seq_lengths_q[-1] == 8 @pytest.mark.skipif(speculator_module.vllm_version_is("0.28.0"), reason="DPSyncState is a main2main interface") diff --git a/vllm_ascend/worker/v2/spec_decode/autoregressive/speculator.py b/vllm_ascend/worker/v2/spec_decode/autoregressive/speculator.py index 8ec7c0c34847..13f0cb1d27e1 100644 --- a/vllm_ascend/worker/v2/spec_decode/autoregressive/speculator.py +++ b/vllm_ascend/worker/v2/spec_decode/autoregressive/speculator.py @@ -496,14 +496,16 @@ def build_draft_attn_metadatas( } if is_draft_model_prefill: - prepared_attn_metadata, _ = self._prepare_replicated_prefill_attn( - attn_metadata, - None, - num_reqs_padded, - num_tokens_padded, - ) - assert prepared_attn_metadata is not None - return [prepared_attn_metadata] + if self.attn_architecture in ("DSA", "SFA"): + prepared_attn_metadata, _ = self._prepare_replicated_prefill_attn( + attn_metadata, + None, + num_reqs_padded, + num_tokens_padded, + ) + assert prepared_attn_metadata is not None + attn_metadata = prepared_attn_metadata + return [attn_metadata] draft_attn_metadatas = self._init_decode_draft_attn_metadatas(attn_metadata, num_reqs_padded)