main
py 58 lines 1.86 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 from vllm.transformers_utils.configs.speculators.base import SpeculatorsConfig
7
8 pytestmark = pytest.mark.skip_global_cleanup
9
10
11 def _make_dspark_config(architecture: str) -> dict:
12 return {
13 "speculators_model_type": "dspark",
14 "architectures": [architecture],
15 "sample_from_anchor": True,
16 "use_aux_hidden_state": True,
17 "draft_vocab_size": 32,
18 "target_hidden_size": 16,
19 "mask_token_id": 31,
20 "markov_rank": 4,
21 "markov_head_type": "vanilla",
22 "block_size": 3,
23 "aux_hidden_state_layer_ids": [2, 5],
24 "transformer_layer_config": {
25 "model_type": "qwen3",
26 "hidden_size": 16,
27 },
28 }
29
30
31 def test_dspark_updater_preserves_qwen3_omni_contract() -> None:
32 config = SpeculatorsConfig.extract_transformers_pre_trained_config(
33 _make_dspark_config("Qwen3OmniDSparkModel")
34 )
35
36 assert config["architectures"] == ["Qwen3OmniDSparkModel"]
37 assert config["sample_from_anchor"] is True
38 assert config["dspark_bonus_anchor"] is False
39 assert config["use_aux_hidden_state"] is True
40 assert config["target_layer_ids"] == [1, 4]
41
42
43 def test_dspark_updater_keeps_legacy_qwen3_checkpoint_loadable() -> None:
44 config = SpeculatorsConfig.extract_transformers_pre_trained_config(
45 _make_dspark_config("DSparkSpeculator")
46 )
47
48 assert config["architectures"] == ["Qwen3DSparkModel"]
49
50
51 def test_dspark_updater_maps_bonus_anchor_semantics() -> None:
52 outer_config = _make_dspark_config("Qwen3DSparkModel")
53 outer_config["sample_from_anchor"] = False
54
55 config = SpeculatorsConfig.extract_transformers_pre_trained_config(outer_config)
56
57 assert config["sample_from_anchor"] is False
58 assert config["dspark_bonus_anchor"] is True