@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 156 lines 5.57 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for ``maybe_override_with_speculators`` config merging.
4
5 When a model is in speculators format, the function previously dropped any
6 user-provided ``--speculative-config`` overrides (e.g. ``attention_backend``
7 added in https://github.com/vllm-project/vllm/pull/39930), silently replacing
8 the dict with a freshly-extracted one. These tests pin the merge behavior so
9 runtime fields from the user are preserved while ``model`` stays auto-detected
10 from the speculators format.
11 """
12
13 from unittest.mock import patch
14
15 from vllm.transformers_utils.config import maybe_override_with_speculators
16
17 _SPECULATORS_CONFIG_DICT = {
18 "architectures": ["DFlashDraftModel"],
19 "speculators_config": {
20 "algorithm": "dflash",
21 "default_proposal_method": "greedy",
22 "proposal_methods": [
23 {
24 "accept_tolerance": 0.0,
25 "proposal_type": "greedy",
26 "speculative_tokens": 8,
27 "verifier_accept_k": 1,
28 }
29 ],
30 "verifier": {
31 "architectures": [],
32 "name_or_path": "fake-org/fake-verifier",
33 },
34 },
35 "speculators_model_type": "dflash",
36 "transformer_layer_config": {},
37 }
38
39
40 def _patched_get_config_dict(config_dict):
41 """Return a patch context that bypasses HF Hub for get_config_dict."""
42 return patch(
43 "vllm.transformers_utils.config.PreTrainedConfig.get_config_dict",
44 return_value=(config_dict, {}),
45 )
46
47
48 def test_speculators_extracts_baseline_config():
49 """Sanity: extracted dict has method/num_speculative_tokens/model fields."""
50 with _patched_get_config_dict(_SPECULATORS_CONFIG_DICT):
51 model, tok, spec = maybe_override_with_speculators(
52 model="fake-org/fake-dflash-spec",
53 tokenizer=None,
54 trust_remote_code=False,
55 vllm_speculative_config=None,
56 )
57
58 assert model == "fake-org/fake-verifier"
59 assert tok == "fake-org/fake-verifier"
60 assert spec is not None
61 assert spec["method"] == "dflash"
62 assert spec["num_speculative_tokens"] == 8
63 assert spec["model"] == "fake-org/fake-dflash-spec"
64
65
66 def test_speculators_preserves_user_attention_backend_override():
67 """User-supplied attention_backend must survive the speculators path.
68
69 Regression for the case where vllm_speculative_config was silently
70 discarded for speculators-format models.
71 """
72 user_overrides = {"attention_backend": "FLASH_ATTN"}
73 with _patched_get_config_dict(_SPECULATORS_CONFIG_DICT):
74 _, _, spec = maybe_override_with_speculators(
75 model="fake-org/fake-dflash-spec",
76 tokenizer=None,
77 trust_remote_code=False,
78 vllm_speculative_config=user_overrides,
79 )
80
81 assert spec is not None
82 assert spec["attention_backend"] == "FLASH_ATTN"
83 # Auto-detected fields still present
84 assert spec["method"] == "dflash"
85 assert spec["num_speculative_tokens"] == 8
86 assert spec["model"] == "fake-org/fake-dflash-spec"
87
88
89 def test_speculators_user_cannot_override_model_field():
90 """``model`` is dictated by the speculators format and must stay locked.
91
92 Even if a user passes a different ``model`` in their override dict, the
93 final config must use the speculators model path (so the drafter actually
94 loads from the right weights).
95 """
96 user_overrides = {"model": "fake-org/wrong-model"}
97 with _patched_get_config_dict(_SPECULATORS_CONFIG_DICT):
98 _, _, spec = maybe_override_with_speculators(
99 model="fake-org/fake-dflash-spec",
100 tokenizer=None,
101 trust_remote_code=False,
102 vllm_speculative_config=user_overrides,
103 )
104
105 assert spec is not None
106 assert spec["model"] == "fake-org/fake-dflash-spec"
107
108
109 def test_speculators_user_cannot_override_method_field():
110 user_overrides = {"method": "ngram"}
111 with _patched_get_config_dict(_SPECULATORS_CONFIG_DICT):
112 _, _, spec = maybe_override_with_speculators(
113 model="fake-org/fake-dflash-spec",
114 tokenizer=None,
115 trust_remote_code=False,
116 vllm_speculative_config=user_overrides,
117 )
118
119 assert spec is not None
120 assert spec["method"] == "dflash"
121
122
123 def test_speculators_user_can_override_num_speculative_tokens():
124 """Runtime knobs (num_speculative_tokens, moe_backend, etc.) take user value."""
125 user_overrides = {
126 "num_speculative_tokens": 4,
127 "moe_backend": "fake-backend",
128 }
129 with _patched_get_config_dict(_SPECULATORS_CONFIG_DICT):
130 _, _, spec = maybe_override_with_speculators(
131 model="fake-org/fake-dflash-spec",
132 tokenizer=None,
133 trust_remote_code=False,
134 vllm_speculative_config=user_overrides,
135 )
136
137 assert spec is not None
138 assert spec["num_speculative_tokens"] == 4
139 assert spec["moe_backend"] == "fake-backend"
140
141
142 def test_non_speculators_model_passes_user_config_through_unchanged():
143 """When the model is NOT a speculators model, user config is returned as-is."""
144 plain_config = {"architectures": ["LlamaForCausalLM"], "hidden_size": 4096}
145 user_overrides = {"method": "ngram", "num_speculative_tokens": 3}
146 with _patched_get_config_dict(plain_config):
147 model, tok, spec = maybe_override_with_speculators(
148 model="fake-org/plain-llama",
149 tokenizer=None,
150 trust_remote_code=False,
151 vllm_speculative_config=user_overrides,
152 )
153
154 assert model == "fake-org/plain-llama"
155 assert tok is None
156 assert spec == user_overrides