diff --git a/tests/ut/patch/platform/test_patch_speculative_config_dspark.py b/tests/ut/patch/platform/test_patch_speculative_config_dspark.py new file mode 100644 index 000000000000..aabc1d694e3f --- /dev/null +++ b/tests/ut/patch/platform/test_patch_speculative_config_dspark.py @@ -0,0 +1,23 @@ +from transformers import Qwen3Config +from vllm.config.speculative import SpeculativeConfig + +import vllm_ascend.patch.platform.patch_speculative_config # noqa: F401 + + +def test_legacy_qwen3_dspark_config_uses_qwen3_loader(): + config = Qwen3Config( + architectures=["DSparkDraftModel"], + block_size=7, + dflash_config={ + "mask_token_id": 163824, + "target_layer_ids": [7, 23, 51, 67, 83], + }, + ) + + normalized = SpeculativeConfig.hf_config_override(config) + + assert normalized is config + assert normalized.architectures == ["Qwen3DSparkModel"] + assert normalized.mask_token_id == 163824 + assert normalized.target_layer_ids == [7, 23, 51, 67, 83] + assert normalized.block_size == 7 diff --git a/tests/ut/spec_decode/test_dspark_proposer.py b/tests/ut/spec_decode/test_dspark_proposer.py index ac486c9aebc9..73857de027f3 100644 --- a/tests/ut/spec_decode/test_dspark_proposer.py +++ b/tests/ut/spec_decode/test_dspark_proposer.py @@ -26,6 +26,7 @@ import numpy as np import pytest import torch +from vllm.v1.worker.utils import AttentionGroup from vllm_ascend.attention.attention_v1 import AscendAttentionState from vllm_ascend.spec_decode.dflash_proposer import AscendDflashProposer @@ -62,7 +63,10 @@ def _make_vllm_config(hf_config: SimpleNamespace) -> SimpleNamespace: """Build the minimal config consumed by the DSpark initializer.""" draft_model_config = SimpleNamespace(hf_config=hf_config, get_hidden_size=lambda: _HIDDEN_SIZE) return SimpleNamespace( - speculative_config=SimpleNamespace(draft_sample_method="greedy", draft_model_config=draft_model_config) + speculative_config=SimpleNamespace( + draft_sample_method="greedy", + draft_model_config=draft_model_config, + ) ) @classmethod @@ -100,7 +104,16 @@ def mock_parent_init( else SimpleNamespace() ) - with patch.object(AscendDSparkProposer.__base__, "__init__", mock_parent_init): + dynamic_spec_config = SimpleNamespace(method="", method_params={}) + with ( + patch.object(AscendDSparkProposer.__base__, "__init__", mock_parent_init), + patch( + "vllm_ascend.spec_decode.dspark_proposer.get_ascend_config", + return_value=SimpleNamespace( + dynamic_spec_config=dynamic_spec_config, + ), + ), + ): proposer = AscendDSparkProposer(vllm_config, device) num_query_total = num_reqs * proposer.num_query_per_req proposer.positions = torch.zeros(max_num_tokens, dtype=torch.int32, device=device) @@ -131,11 +144,11 @@ def mock_parent_init( proposer._per_group_block_table_buffers = {gid: block_table} slot = torch.zeros(max_num_tokens, dtype=torch.int32, device=device) proposer._per_group_slot_mappings = {gid: slot} + proposer._per_group_kernel_block_sizes = {gid: block_size} proposer._per_group_query_slot_mapping_buffers = {gid: slot.clone()} proposer._per_group_context_slot_mapping_buffers = {gid: slot.clone()} return proposer - # fmt: off @staticmethod def _invoke_set_inputs_first_pass( proposer, @@ -143,6 +156,8 @@ def _invoke_set_inputs_first_pass( num_reqs, block_size, seq_len=128, + host_seq_len=None, + async_metadata=False, context=None, num_rejected=None, with_optional_attrs=False, @@ -155,17 +170,20 @@ def _invoke_set_inputs_first_pass( next_token_ids, target_hidden_states)``. """ next_token_ids = torch.arange(1, num_reqs + 1, dtype=torch.int64) - target_hidden_states = torch.arange( - num_reqs * 8, dtype=torch.float32 - ).reshape(num_reqs, 8) + target_hidden_states = torch.arange(num_reqs * 8, dtype=torch.float32).reshape(num_reqs, 8) query_start_loc_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32) if context is not None: query_start_loc_cpu[num_reqs] = context + if host_seq_len is None: + host_seq_len = seq_len + seq_lens_cpu = torch.full((num_reqs,), host_seq_len, dtype=torch.int32) cad = SimpleNamespace( num_reqs=num_reqs, query_start_loc=torch.arange(num_reqs + 1, dtype=torch.int32) * block_size, query_start_loc_cpu=query_start_loc_cpu, seq_lens=torch.full((num_reqs,), seq_len, dtype=torch.int32), + _seq_lens_cpu=seq_lens_cpu, + seq_lens_cpu=None if async_metadata else seq_lens_cpu, max_seq_len=seq_len, ) if with_optional_attrs: @@ -183,9 +201,6 @@ def _invoke_set_inputs_first_pass( return num_query_total, token_indices, cad, extra, next_token_ids, target_hidden_states -# fmt: on - - class TestDSparkPositionsFullUnderMultiDp(_DSparkProposerTestBase): """Guard: under multi-DP the dspark draft proposer must hand DSA attention a full-length positions buffer so ``positions[:num_input_tokens]`` never reads @@ -368,7 +383,6 @@ def test_configures_anchor_sampling( assert proposer.max_query_tokens == expected_max_query_tokens -# fmt: off class TestSetPerGroupAttnMetadata(_DSparkProposerTestBase): """``set_per_group_attn_metadata`` stores the runner-provided per-group block table / slot mapping into the read-only dicts the proposer consults @@ -376,9 +390,7 @@ class TestSetPerGroupAttnMetadata(_DSparkProposerTestBase): def test_stores_block_table_and_slot_mapping(self): num_reqs, block_size, max_num_tokens = 4, 5, 256 - proposer = self._make_proposer( - max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size - ) + proposer = self._make_proposer(max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size) # a gid not pre-populated by _make_proposer (which only seeds gid=0) gid = 7 block_table = torch.zeros((num_reqs, 16), dtype=torch.int32) @@ -391,9 +403,7 @@ def test_stores_block_table_and_slot_mapping(self): def test_overwrites_existing_gid(self): num_reqs, block_size, max_num_tokens = 2, 5, 256 - proposer = self._make_proposer( - max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size - ) + proposer = self._make_proposer(max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size) gid = 0 # already populated by _make_proposer old_block_table = proposer._per_group_block_tables[gid] new_block_table = torch.ones((num_reqs, 16), dtype=torch.int32) @@ -424,10 +434,7 @@ def _make_vllm_config( speculative_config = SimpleNamespace( num_speculative_tokens=num_speculative_tokens, draft_sample_method=draft_sample_method, - draft_model_config=SimpleNamespace( - hf_config=SimpleNamespace(), - get_hidden_size=lambda: hidden_size - ), + draft_model_config=SimpleNamespace(hf_config=SimpleNamespace(), get_hidden_size=lambda: hidden_size), ) return SimpleNamespace(speculative_config=speculative_config) @@ -451,7 +458,6 @@ def _stub(self, vllm_config, device, runner=None): self.dtype = dtype self.device = device self.draft_model_config = vllm_config.speculative_config.draft_model_config - # present so the ``del`` in DSpark.__init__ succeeds self.hidden_size = 0 self.hidden_states = None self._dflash_hidden_states = None @@ -498,7 +504,14 @@ def test_greedy_allocates_dspark_buffers(self, monkeypatch): draft_sample_method="greedy", hidden_size=hidden, ) - proposer = AscendDSparkProposer(vllm_config, device) + dynamic_spec_config = SimpleNamespace(method="", method_params={}) + with patch( + "vllm_ascend.spec_decode.dspark_proposer.get_ascend_config", + return_value=SimpleNamespace( + dynamic_spec_config=dynamic_spec_config, + ), + ): + proposer = AscendDSparkProposer(vllm_config, device) blk = 1 + num_spec max_query_tokens = max_batch * num_spec @@ -533,21 +546,16 @@ class TestSetInputsFirstPassOutputs(_DSparkProposerTestBase): @pytest.fixture(autouse=True) def _mock_kernel(self, monkeypatch): monkeypatch.setattr( - "vllm_ascend.spec_decode.dspark_proposer." - "copy_and_expand_dflash_and_dspark_inputs_kernel", + "vllm_ascend.spec_decode.dspark_proposer.copy_and_expand_dflash_and_dspark_inputs_kernel", MagicMock(), ) def test_return_value_and_token_indices(self): num_reqs, block_size, max_num_tokens = 4, 5, 256 - proposer = self._make_proposer( - max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size - ) - num_query_total, token_indices, _cad, extra = ( - self._invoke_set_inputs_first_pass( - proposer, num_reqs=num_reqs, block_size=block_size - )[:4] - ) + proposer = self._make_proposer(max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size) + num_query_total, token_indices, _cad, extra = self._invoke_set_inputs_first_pass( + proposer, num_reqs=num_reqs, block_size=block_size + )[:4] assert num_query_total == num_reqs * block_size assert token_indices.shape == (num_reqs * block_size,) assert token_indices.dtype == torch.int32 @@ -556,33 +564,49 @@ def test_return_value_and_token_indices(self): def test_seed_buffer_copied_from_next_tokens(self): num_reqs, block_size, max_num_tokens = 4, 5, 256 - proposer = self._make_proposer( - max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size - ) - self._invoke_set_inputs_first_pass( - proposer, num_reqs=num_reqs, block_size=block_size - ) + proposer = self._make_proposer(max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size) + self._invoke_set_inputs_first_pass(proposer, num_reqs=num_reqs, block_size=block_size) expected = torch.arange(1, num_reqs + 1, dtype=torch.int64) assert torch.equal(proposer._dspark_seed_buffer[:num_reqs], expected) assert torch.all(proposer._dspark_seed_buffer[num_reqs:] == 0) def test_context_hidden_states_copied(self): num_reqs, block_size, max_num_tokens = 4, 5, 256 + proposer = self._make_proposer(max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size) + self._invoke_set_inputs_first_pass(proposer, num_reqs=num_reqs, block_size=block_size, context=num_reqs) + assert proposer._dflash_num_context == num_reqs + expected = torch.arange(num_reqs * 8, dtype=torch.float32).reshape(num_reqs, 8) + assert torch.equal(proposer._dflash_hidden_states[:num_reqs], expected) + + def test_query_slot_kernel_uses_logical_block_size(self, monkeypatch): + kernel = MagicMock() + monkeypatch.setattr( + "vllm_ascend.spec_decode.dspark_proposer.copy_and_expand_dflash_and_dspark_inputs_kernel", + kernel, + ) + num_reqs, num_speculative_tokens, max_num_tokens = 1, 7, 32 proposer = self._make_proposer( - max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size + max_num_tokens=max_num_tokens, + num_reqs=num_reqs, + block_size=num_speculative_tokens, ) + proposer.draft_attn_groups[0].kv_cache_spec.block_size = 384 + proposer._per_group_kernel_block_sizes[0] = 128 + self._invoke_set_inputs_first_pass( - proposer, num_reqs=num_reqs, block_size=block_size, context=num_reqs + proposer, + num_reqs=num_reqs, + block_size=num_speculative_tokens, + seq_len=720, ) - assert proposer._dflash_num_context == num_reqs - expected = torch.arange(num_reqs * 8, dtype=torch.float32).reshape(num_reqs, 8) - assert torch.equal(proposer._dflash_hidden_states[:num_reqs], expected) + + kwargs = kernel[1,].call_args.kwargs + assert proposer.draft_attn_groups[0].kv_cache_spec.block_size == 384 + assert kwargs["block_size"] == 128 def test_cad_rewritten_to_cross_attention_shape(self): num_reqs, block_size, max_num_tokens = 4, 5, 256 - proposer = self._make_proposer( - max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size - ) + proposer = self._make_proposer(max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size) num_query_total, _, cad, _ = self._invoke_set_inputs_first_pass( proposer, num_reqs=num_reqs, block_size=block_size, with_optional_attrs=True )[:4] @@ -600,10 +624,7 @@ def test_cad_rewritten_to_cross_attention_shape(self): # slot mapping is a slice of the primary group's query buffer (shares # storage from offset 0); a fresh slice is not identity-equal, so check # the underlying storage and length instead. - assert ( - cad.slot_mapping.data_ptr() - == proposer._per_group_query_slot_mapping_buffers[0].data_ptr() - ) + assert cad.slot_mapping.data_ptr() == proposer._per_group_query_slot_mapping_buffers[0].data_ptr() assert cad.slot_mapping.shape[0] == num_query_total # optional attrs the proposer rewrites when present. assert cad.actual_seq_lengths_q == [block_size] * num_reqs @@ -617,25 +638,24 @@ def test_cad_uses_model_reported_causality(self): block_size=block_size, draft_attn_causal=True, ) - _, _, cad, _ = self._invoke_set_inputs_first_pass( - proposer, num_reqs=num_reqs, block_size=block_size - )[:4] + _, _, cad, _ = self._invoke_set_inputs_first_pass(proposer, num_reqs=num_reqs, block_size=block_size)[:4] assert cad.causal is True def test_cad_query_start_loc_and_seq_lens(self): num_reqs, block_size, max_num_tokens = 4, 5, 256 - proposer = self._make_proposer( - max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size - ) - _nqt, _ti, cad, _extra = self._invoke_set_inputs_first_pass( - proposer, num_reqs=num_reqs, block_size=block_size - )[:4] + proposer = self._make_proposer(max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size) + _nqt, _ti, cad, _extra = self._invoke_set_inputs_first_pass(proposer, num_reqs=num_reqs, block_size=block_size)[ + :4 + ] expected_qsl = torch.arange(num_reqs + 1, dtype=torch.int32) * block_size assert torch.equal(cad.query_start_loc, expected_qsl) assert torch.equal(cad.query_start_loc_cpu, expected_qsl) # seq_lens grow by block_size when no tokens were rejected. - assert torch.equal(cad.seq_lens, torch.full((num_reqs,), 128 + block_size, dtype=torch.int32)) + expected = torch.full((num_reqs,), 128 + block_size, dtype=torch.int32) + assert torch.equal(cad.seq_lens, expected) + assert torch.equal(cad._seq_lens_cpu, expected) + assert torch.equal(cad.seq_lens_cpu, expected) class TestSetInputsFirstPassRejectedTokens(_DSparkProposerTestBase): @@ -644,38 +664,36 @@ class TestSetInputsFirstPassRejectedTokens(_DSparkProposerTestBase): def test_seq_lens_subtracts_rejected(self, monkeypatch): monkeypatch.setattr( - "vllm_ascend.spec_decode.dspark_proposer." - "copy_and_expand_dflash_and_dspark_inputs_kernel", + "vllm_ascend.spec_decode.dspark_proposer.copy_and_expand_dflash_and_dspark_inputs_kernel", MagicMock(), ) num_reqs, block_size, max_num_tokens = 4, 5, 256 - proposer = self._make_proposer( - max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size - ) + proposer = self._make_proposer(max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size) rejected = torch.full((num_reqs,), 2, dtype=torch.int32) _nqt, _ti, cad, _extra = self._invoke_set_inputs_first_pass( - proposer, num_reqs=num_reqs, block_size=block_size, num_rejected=rejected + proposer, + num_reqs=num_reqs, + block_size=block_size, + host_seq_len=126, + async_metadata=True, + num_rejected=rejected, )[:4] # effective = seq_lens(128) - rejected(2) = 126; then + block_size(5) = 131. - assert torch.equal( - cad.seq_lens, torch.full((num_reqs,), 128 - 2 + block_size, dtype=torch.int32) - ) + assert torch.equal(cad.seq_lens, torch.full((num_reqs,), 128 - 2 + block_size, dtype=torch.int32)) + expected_host = torch.full((num_reqs,), 126 + block_size, dtype=torch.int32) + assert torch.equal(cad._seq_lens_cpu, expected_host) + assert cad.seq_lens_cpu is None def test_kernel_called_with_has_num_rejected(self, monkeypatch): kernel = MagicMock() monkeypatch.setattr( - "vllm_ascend.spec_decode.dspark_proposer." - "copy_and_expand_dflash_and_dspark_inputs_kernel", + "vllm_ascend.spec_decode.dspark_proposer.copy_and_expand_dflash_and_dspark_inputs_kernel", kernel, ) num_reqs, block_size, max_num_tokens = 4, 5, 256 - proposer = self._make_proposer( - max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size - ) + proposer = self._make_proposer(max_num_tokens=max_num_tokens, num_reqs=num_reqs, block_size=block_size) rejected = torch.full((num_reqs,), 2, dtype=torch.int32) - self._invoke_set_inputs_first_pass( - proposer, num_reqs=num_reqs, block_size=block_size, num_rejected=rejected - ) + self._invoke_set_inputs_first_pass(proposer, num_reqs=num_reqs, block_size=block_size, num_rejected=rejected) # The proposer calls the kernel as ``kernel[1,](...)`` (Triton-style # grid indexing), so the call lands on the indexed sub-mock. sub = kernel[1,] @@ -686,10 +704,8 @@ def test_kernel_called_with_has_num_rejected(self, monkeypatch): assert kwargs["SAMPLE_FROM_ANCHOR"] is True -class TestInitializeAttnBackendErrors(_DSparkProposerTestBase): - """``initialize_attn_backend`` raises clearly when the draft model does not - expose the DSpark layer-name API, or when no draft attention groups can be - built from the kv-cache groups.""" +class TestInitializeAttnBackend(_DSparkProposerTestBase): + """Initialization preserves each group's logical kernel block size.""" @staticmethod def _make_proposer_for_init(): @@ -698,32 +714,45 @@ def _make_proposer_for_init(): proposer.device = torch.device("cpu") return proposer - def test_model_without_draft_layer_names_raises(self, monkeypatch): - # get_layers_from_vllm_config is called first; stub it so the model - # check is what actually fails. + def test_initialization_tracks_logical_block_size_per_gid(self, monkeypatch): + manager_specs = [MagicMock(), MagicMock()] + for spec in manager_specs: + spec.block_size = 384 + + backend = MagicMock() + backend.full_cls_name.return_value = "fake.backend" + layers = {} + for gid in range(2): + layer = MagicMock() + layer.get_attn_backend.return_value = backend + layers[f"L{gid}"] = layer monkeypatch.setattr( "vllm_ascend.spec_decode.dspark_proposer.get_layers_from_vllm_config", - lambda *a, **k: {}, + lambda *a, **k: layers, ) - proposer = self._make_proposer_for_init() - # model lacks get_draft_kv_cache_layer_names entirely. - proposer.model = SimpleNamespace() - kv_cache_config = SimpleNamespace(kv_cache_groups=[]) - with pytest.raises(RuntimeError, match="get_draft_kv_cache_layer_names"): - proposer.initialize_attn_backend(kv_cache_config) - - def test_no_draft_attn_groups_raises(self, monkeypatch): - monkeypatch.setattr( - "vllm_ascend.spec_decode.dspark_proposer.get_layers_from_vllm_config", - lambda *a, **k: {}, - ) proposer = self._make_proposer_for_init() - # draft layer names exist, but no kv-cache group names overlap them. - proposer.model = SimpleNamespace(get_draft_kv_cache_layer_names=lambda: {"L0"}) - - non_overlapping_group = SimpleNamespace(layer_names=["OTHER_LAYER"]) - kv_cache_config = SimpleNamespace(kv_cache_groups=[non_overlapping_group]) - with pytest.raises(RuntimeError, match="registered draft attention groups"): - proposer.initialize_attn_backend(kv_cache_config) -# fmt: on + proposer.model = SimpleNamespace(get_draft_kv_cache_layer_names=lambda: {"L0", "L1"}) + proposer.max_query_tokens = 8 + proposer.max_num_tokens = 16 + kv_cache_config = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace( + layer_names=[f"L{gid}"], + kv_cache_spec=manager_specs[gid], + ) + for gid in range(2) + ], + ) + + with patch.object(AttentionGroup, "create_metadata_builders") as create_builders: + proposer.initialize_attn_backend( + kv_cache_config, + kernel_block_sizes=[128, 64], + ) + + assert [spec.block_size for spec in manager_specs] == [384, 384] + assert proposer._per_group_kernel_block_sizes == {0: 128, 1: 64} + assert [group.kv_cache_group_id for group in proposer.draft_attn_groups] == [0, 1] + assert proposer.kernel_block_size == 128 + assert [call.kwargs["kernel_block_size"] for call in create_builders.call_args_list] == [128, 64] diff --git a/tests/ut/spec_decode/test_llm_base_proposer.py b/tests/ut/spec_decode/test_llm_base_proposer.py index 0f78198a8ac2..2c8be1bfcd38 100644 --- a/tests/ut/spec_decode/test_llm_base_proposer.py +++ b/tests/ut/spec_decode/test_llm_base_proposer.py @@ -76,16 +76,22 @@ def test_pixtral_uses_vision_config_image_token_id(self): assert image_token_index == 789 - def test_kimi_uses_media_placeholder_token_id(self): + @pytest.mark.parametrize( + "model_name", + [ + "KimiK25ForConditionalGeneration", + "KimiK3ForConditionalGeneration", + "AscendKimiK3ForConditionalGeneration", + ], + ) + def test_kimi_uses_media_placeholder_token_id(self, model_name: str): config = SimpleNamespace( image_token_id=123, image_token_index=456, media_placeholder_token_id=789, ) - image_token_index = AscendSpecDecodeBaseProposer._get_multimodal_image_token_index( - "KimiK25ForConditionalGeneration", config - ) + image_token_index = AscendSpecDecodeBaseProposer._get_multimodal_image_token_index(model_name, config) assert image_token_index == 789 @@ -103,7 +109,8 @@ def test_load_model_reads_validated_draft_window_size(): proposer = AscendSpecDecodeBaseProposer.__new__(AscendSpecDecodeBaseProposer) proposer.vllm_config = SimpleNamespace(additional_config={"draft_window_size": 64}) proposer.maybe_eager_context = nullcontext() - proposer._get_model = MagicMock(return_value=MagicMock()) + draft_model = MagicMock() + proposer._get_model = MagicMock(return_value=draft_model) proposer.method = "eagle3" proposer.num_speculative_tokens = 4 proposer.runner = SimpleNamespace(max_num_reqs=8) diff --git a/tests/ut/spec_decode/test_utils.py b/tests/ut/spec_decode/test_utils.py index 494082f4bd3f..575d53cb89c7 100644 --- a/tests/ut/spec_decode/test_utils.py +++ b/tests/ut/spec_decode/test_utils.py @@ -8,10 +8,13 @@ from __future__ import annotations +from types import SimpleNamespace + import numpy as np import torch from vllm_ascend.spec_decode.utils import ( + SlidingWindowAdapter, correct_optimistic_seq_lens_cpu, update_num_computed_tokens_for_batch_change, ) @@ -184,3 +187,28 @@ def test_cpu_and_gpu_corrections_agree(): gpu_seq_lens = num_computed_gpu.numpy() + num_scheduled_step_n np.testing.assert_array_equal(optimistic, gpu_seq_lens) + + +def test_sliding_window_updates_host_length_mirrors(): + adapter = SlidingWindowAdapter( + window_size=128, + block_size=64, + max_num_reqs=1, + future_offset=0, + device=torch.device("cpu"), + ) + metadata = SimpleNamespace( + block_table_tensor=torch.arange(10, dtype=torch.int32).view(1, 10), + seq_lens=torch.tensor([300], dtype=torch.int32), + seq_lens_cpu=torch.tensor([300], dtype=torch.int32), + _seq_lens_cpu=torch.tensor([300], dtype=torch.int32), + seq_lens_cpu_upper_bound=torch.tensor([300], dtype=torch.int32), + ) + + adapter.apply(metadata) + + expected = torch.tensor([172], dtype=torch.int32) + torch.testing.assert_close(metadata.seq_lens, expected) + torch.testing.assert_close(metadata.seq_lens_cpu, expected) + torch.testing.assert_close(metadata._seq_lens_cpu, expected) + torch.testing.assert_close(metadata.seq_lens_cpu_upper_bound, expected) diff --git a/vllm_ascend/patch/platform/patch_speculative_config.py b/vllm_ascend/patch/platform/patch_speculative_config.py index c4f4737c8043..269607a336de 100644 --- a/vllm_ascend/patch/platform/patch_speculative_config.py +++ b/vllm_ascend/patch/platform/patch_speculative_config.py @@ -1,6 +1,23 @@ +from transformers import PretrainedConfig from vllm.config.speculative import SpeculativeConfig _orig_post_init = SpeculativeConfig.__post_init__ +_orig_hf_config_override = SpeculativeConfig.hf_config_override + + +def _normalize_legacy_qwen3_dspark_config(hf_config: PretrainedConfig) -> PretrainedConfig: + hf_config = _orig_hf_config_override(hf_config) + architectures = hf_config.architectures or () + if hf_config.model_type == "qwen3" and "DSparkDraftModel" in architectures: + dflash_config = hf_config.dflash_config + hf_config.update( + { + "architectures": ["Qwen3DSparkModel"], + "mask_token_id": dflash_config["mask_token_id"], + "target_layer_ids": dflash_config["target_layer_ids"], + } + ) + return hf_config def _dspark_post_init(self): @@ -16,4 +33,5 @@ def _dspark_post_init(self): draft_hf_config.ptd_token_id = getattr(draft_hf_config, "mask_token_id", None) # type: ignore +SpeculativeConfig.hf_config_override = staticmethod(_normalize_legacy_qwen3_dspark_config) SpeculativeConfig.__post_init__ = _dspark_post_init diff --git a/vllm_ascend/spec_decode/dspark_proposer.py b/vllm_ascend/spec_decode/dspark_proposer.py index ac1485c8268f..643f2128aae1 100644 --- a/vllm_ascend/spec_decode/dspark_proposer.py +++ b/vllm_ascend/spec_decode/dspark_proposer.py @@ -43,8 +43,8 @@ def __init__( blk = 1 + self.num_speculative_tokens self._dspark_draft_buffer = torch.zeros((self.max_batch_size, blk), dtype=torch.int64, device=device) self._dspark_seed_buffer = torch.zeros(self.max_batch_size, dtype=torch.int64, device=device) - # DSpark is not supported in vllm v1, so related property needs to be reset here. - del self.hidden_size, self.hidden_states, self._dflash_hidden_states # type: ignore[has-type] + # Replace the target-sized DFlash buffers with the draft model's hidden + # size. Assignment releases the old tensors without an explicit del. self.hidden_size = vllm_config.speculative_config.draft_model_config.get_hidden_size() self.hidden_states = torch.zeros( (self.max_num_tokens, self.hidden_size), @@ -87,27 +87,20 @@ def __init__( device=device, ) - # TODO simplify these comments - # block_table / slot_mapping bookkeeping (10 dicts below). v1 self- - # manages per kv_cache_group_id / per layer because it lacks v2's - # BlockTables scaffold; v2 injects a single self.block_tables - # (BlockTables, with .slot_mappings) + build_slot_mappings_by_layer, - # so the speculator holds none of these. P2 refactor target (move to - # runner). - - # per-gid block_table from runner (just read) + # The v1 runner owns block tables and slot mappings. Keep per-group + # references here because K3 draft layers can span multiple cache + # groups with different logical block sizes. self._per_group_block_tables: dict[int, torch.Tensor] = {} - # per-gid slot_mapping from runner (just read) self._per_group_slot_mappings: dict[int, torch.Tensor] = {} + # Per-gid logical block size used to expand slot mappings. The KV + # manager's physical page can be larger when hybrid cache groups share + # one allocation, so kv_cache_spec.block_size is not interchangeable + # with the attention kernel's block size. + self._per_group_kernel_block_sizes: dict[int, int] = {} - # per-gid block_table (use in proposer) self._per_group_block_table_buffers: dict[int, torch.Tensor] = {} - # per-gid query slot_mapping buffer self._per_group_query_slot_mapping_buffers: dict[int, torch.Tensor] = {} - # per-gid context slot_mapping buffer self._per_group_context_slot_mapping_buffers: dict[int, torch.Tensor] = {} - - # per-layer context slot mappings as a flat list self._context_slot_mapping_buffers: list[torch.Tensor | None] | None = None def _compute_confidence( @@ -129,24 +122,22 @@ def _compute_confidence( confidence.copy_(conf_raw.reshape(num_reqs, self.num_speculative_tokens)) return confidence - def initialize_attn_backend(self, kv_cache_config, kernel_block_sizes=None) -> None: + def initialize_attn_backend( + self, + kv_cache_config, + kernel_block_sizes: list[int] | None = None, + ) -> None: # Find draft layers (attention layers added by draft model) all_attn_layers = get_layers_from_vllm_config( self.vllm_config, AttentionLayerBase, # type: ignore[type-abstract] ) - attention_groups_list: list[dict[tuple[str, str], AttentionGroup]] = [] - # the draft layers have multiple kv_cache_groups - if not hasattr(self.model, "get_draft_kv_cache_layer_names"): - raise RuntimeError( - "DSpark standard-cache path requires the draft model to expose get_draft_kv_cache_layer_names" - ) - self._draft_attn_layer_names = set(self.model.get_draft_kv_cache_layer_names()) self.attn_layer_names = list(sorted(self._draft_attn_layer_names)) + self._per_group_kernel_block_sizes = {} + self.draft_attn_groups: list[AttentionGroup] = [] - # there are many kv groups other than one for kv_cache_gid, kv_cache_group_spec in enumerate(kv_cache_config.kv_cache_groups): draft_layer_names_in_group = set(kv_cache_group_spec.layer_names) & self._draft_attn_layer_names if not draft_layer_names_in_group: @@ -162,33 +153,31 @@ def initialize_attn_backend(self, kv_cache_config, kernel_block_sizes=None) -> N key = (attn_backend.full_cls_name(), layer_kv_cache_spec) if key not in attention_groups: + kernel_block_size = int( + kernel_block_sizes[kv_cache_gid] + if kernel_block_sizes is not None and kv_cache_gid < len(kernel_block_sizes) + else layer_kv_cache_spec.block_size + ) attn_group = AttentionGroup( attn_backend, [layer_name], layer_kv_cache_spec, kv_cache_gid, ) - attn_group.create_metadata_builders(self.vllm_config, self.device) + attn_group.create_metadata_builders( + self.vllm_config, + self.device, + kernel_block_size=kernel_block_size, + ) + self._per_group_kernel_block_sizes[kv_cache_gid] = kernel_block_size attention_groups[key] = attn_group else: attention_groups[key].layer_names.append(layer_name) - attention_groups_list.append(attention_groups) - - self.draft_attn_groups = [ - attention_group - for attention_groups in attention_groups_list - for attention_group in attention_groups.values() - ] - self.kv_cache_gid = 0 - if not self.draft_attn_groups: - raise RuntimeError( - "DSpark standard-cache path requires registered draft attention " - f"groups. Missing layers: {self.attn_layer_names}" - ) + self.draft_attn_groups.extend(attention_groups.values()) self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id - self.kernel_block_size = int(self.draft_attn_groups[0].kv_cache_spec.block_size) + self.kernel_block_size = self._per_group_kernel_block_sizes[self.kv_cache_gid] name_to_gid = { ln: gid @@ -257,13 +246,10 @@ def set_inputs_first_pass( # Query block: reuse the DFlash inputs kernel logic (host-side ref) # per kv-cache-group to fill positions / input_ids / query slot_mapping # / token_indices. - draft_attn_groups = getattr(self, "draft_attn_groups", []) - for attn_group in draft_attn_groups: + for attn_group in self.draft_attn_groups: gid = attn_group.kv_cache_group_id - gid_block_table = self._per_group_block_table_buffers.get(gid) - if gid_block_table is None: - continue - kv_block_size = int(attn_group.kv_cache_spec.block_size) + gid_block_table = self._per_group_block_table_buffers[gid] + kernel_block_size = self._per_group_kernel_block_sizes[gid] copy_and_expand_dflash_and_dspark_inputs_kernel[ (_compute_num_programs(self._dflash_num_context, num_query_total),) ]( @@ -287,7 +273,7 @@ def set_inputs_first_pass( num_rejected_tokens_ptr=num_rejected_tokens_gpu, # Scalars parallel_drafting_token_id=self.parallel_drafting_token_id, - block_size=kv_block_size, + block_size=kernel_block_size, num_query_per_req=self.num_query_per_req, num_speculative_tokens=self.num_speculative_tokens, total_input_tokens=self._dflash_num_context, @@ -306,6 +292,15 @@ def set_inputs_first_pass( cad.query_start_loc = self.arange_dflash[: batch_size + 1] * self.num_query_per_req cad.seq_lens = effective_seq_lens + self.num_query_per_req + # The model runner has already corrected this canonical host mirror + # with the accepted-token count. Extend it on CPU alongside the device + # lengths, without another reject D2H copy or attention-side wait. + if cad._seq_lens_cpu is not None: + draft_seq_lens_cpu = cad._seq_lens_cpu.clone() + draft_seq_lens_cpu[:batch_size].add_(self.num_query_per_req) + cad._seq_lens_cpu = draft_seq_lens_cpu + if getattr(cad, "seq_lens_cpu", None) is not None: + cad.seq_lens_cpu = draft_seq_lens_cpu cad.query_start_loc_cpu = ( torch.from_numpy(self.token_arange_np[: batch_size + 1]).clone() * self.num_query_per_req ).to(torch.int32) diff --git a/vllm_ascend/spec_decode/llm_base_proposer.py b/vllm_ascend/spec_decode/llm_base_proposer.py index 8c5121e20a6d..d6bfd42f7131 100644 --- a/vllm_ascend/spec_decode/llm_base_proposer.py +++ b/vllm_ascend/spec_decode/llm_base_proposer.py @@ -26,6 +26,7 @@ from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM +from vllm.models.kimi_k3.nvidia.dspark_mla import K3DSparkForCausalLM from vllm.triton_utils import HAS_TRITON, triton from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.attention.backends.utils import CommonAttentionMetadata @@ -66,6 +67,16 @@ # Currently we will fix block size to a small one since `num_reqs` can't be too large _PREPARE_INPUTS_BLOCK_SIZE = 4 +_HIDDEN_STATE_DRAFTER_TYPES = ( + Eagle3LlamaForCausalLM, + DFlashQwen3ForCausalLM, + Qwen3DSparkForCausalLM, + K3DSparkForCausalLM, + Eagle3VwnLlamaForCausalLM, + Eagle3DeepseekV2ForCausalLM, + DSparkDeepseekV4ForCausalLM, +) + def greedy_sample(logits: torch.Tensor) -> torch.Tensor: tp_group = get_tp_group() @@ -115,7 +126,11 @@ def _get_multimodal_image_token_index(model_name: str, config: Any) -> int: return config.image_token_id if model_name == "PixtralForConditionalGeneration": return config.vision_config.image_token_id - if model_name == "KimiK25ForConditionalGeneration": + if model_name in { + "KimiK25ForConditionalGeneration", + "KimiK3ForConditionalGeneration", + "AscendKimiK3ForConditionalGeneration", + }: return config.media_placeholder_token_id return config.image_token_index @@ -484,17 +499,20 @@ def _maybe_share_lm_head(self, model: nn.Module) -> None: ) else: logger.info("[spec_decode/base] Loading EAGLE/DFLASH LM head weights from the target model.") + target_lm_head = None if hasattr(model, "lm_head"): - self.model.lm_head = model.lm_head + target_lm_head = model.lm_head elif hasattr(model, "get_language_model") and hasattr(model.get_language_model(), "lm_head"): - self.model.lm_head = model.get_language_model().lm_head - else: + target_lm_head = model.get_language_model().lm_head + if target_lm_head is None: logger.warning( "[spec_decode/base] Target model has no accessible lm_head" " for sharing. Draft model will use its own lm_head." " This may cause incorrect logits if the draft lm_head" " is not trained." ) + else: + self.model.lm_head = target_lm_head if self.method == "mtp" and self.vllm_config.model_config.is_deepseek_mla: for _, layer_module in self.model.model.layers.items(): @@ -793,18 +811,8 @@ def _propose( model = self.model if isinstance(model, BreakableACLGraphWrapper): model = model.unwrap() - assert isinstance( - model, - ( - Eagle3LlamaForCausalLM, - DFlashQwen3ForCausalLM, - Qwen3DSparkForCausalLM, - Eagle3VwnLlamaForCausalLM, - Eagle3DeepseekV2ForCausalLM, - DSparkDeepseekV4ForCausalLM, - ), - ) - target_hidden_states = self.model.combine_hidden_states(target_hidden_states) + assert isinstance(model, _HIDDEN_STATE_DRAFTER_TYPES) + target_hidden_states = model.combine_hidden_states(target_hidden_states) assert target_hidden_states.shape[-1] == self.hidden_size num_tokens, token_indices_to_sample, common_attn_metadata, long_seq_args = self.set_inputs_first_pass( @@ -878,6 +886,19 @@ def _propose( ) if self.method == "dflash": common_attn_metadata.seq_lens = self._adjust_tensor(common_attn_metadata.seq_lens, num_reqs_padded) + elif self.method == "dspark": + # DSpark already rewrote both device and host sequence lengths + # in set_inputs_first_pass. Preserve those values while + # extending only the padded tail for full-graph replay. + common_attn_metadata.seq_lens = self._adjust_tensor(common_attn_metadata.seq_lens, num_reqs_padded) + if common_attn_metadata.seq_lens_cpu is not None: + common_attn_metadata.seq_lens_cpu = self._adjust_tensor( + common_attn_metadata.seq_lens_cpu, num_reqs_padded + ) + if common_attn_metadata._seq_lens_cpu is not None: + common_attn_metadata._seq_lens_cpu = self._adjust_tensor( + common_attn_metadata._seq_lens_cpu, num_reqs_padded + ) else: common_attn_metadata.seq_lens = self._adjust_tensor(self.runner.seq_lens, num_reqs_padded) common_attn_metadata.seq_lens_cpu = self._adjust_tensor(