@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 209 lines 9.42 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 SUPPORTED_SPECULATORS_TYPES = {}
5
6
7 def register_speculator(name):
8 def decorator(fn):
9 SUPPORTED_SPECULATORS_TYPES[name] = fn
10 return fn
11
12 return decorator
13
14
15 @register_speculator("eagle3")
16 def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None:
17 """Apply Eagle-3 specific configuration transformations to the `dict` used to
18 construct the Transformers PreTrainedConfig.
19
20 Eagle-3 specific fields:
21 - draft_vocab_size: Size of the draft model's vocabulary
22 - target_hidden_size: Hidden size of the target model
23 - norm_before_residual: Whether to apply norm before residual connection
24 - norm_before_fc: Whether to apply RMSNorm before the fc projection
25 - eagle_aux_hidden_state_layer_ids: List of layer indices from the base
26 model to use as auxiliary inputs for the Eagle3 drafter. These layers
27 provide intermediate hidden states that help the drafter make better
28 predictions. This is the standard field used in Eagle3 checkpoints.
29 """
30 pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
31 if config_dict.get("target_hidden_size") is not None:
32 pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"]
33 pre_trained_config["norm_before_residual"] = config_dict.get(
34 "norm_before_residual", True
35 )
36 pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False)
37 pre_trained_config["fc_norm"] = config_dict.get("fc_norm", False)
38 pre_trained_config["norm_output"] = config_dict.get("norm_output", False)
39 eagle3_arch_map = {
40 "qwen3": "Eagle3Qwen3ForCausalLM",
41 "llama": "Eagle3LlamaForCausalLM",
42 }
43 model_type = pre_trained_config.get("model_type", "llama")
44 if model_type not in eagle3_arch_map:
45 raise ValueError(f"Unsupported model_type {model_type} for Eagle3 speculator")
46 pre_trained_config["architectures"] = [eagle3_arch_map[model_type]]
47 if config_dict.get("eagle_aux_hidden_state_layer_ids"):
48 pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[
49 "eagle_aux_hidden_state_layer_ids"
50 ]
51
52
53 @register_speculator("peagle")
54 def update_peagle(config_dict: dict, pre_trained_config: dict) -> None:
55 """Apply PEagle (Parallel Eagle) specific configuration transformations to
56 the `dict` used to construct the Transformers PreTrainedConfig.
57
58 PEagle specific fields:
59 - draft_vocab_size: Size of the draft model's vocabulary
60 - target_hidden_size: Hidden size of the target model
61 - norm_before_residual: Whether to apply norm before residual connection
62 - norm_before_fc: Whether to apply RMSNorm before the fc projection
63 - mask_token_id (required): Token ID used for parallel drafting mask
64 placeholders, mapped to pard_token for the proposer
65 - eagle_aux_hidden_state_layer_ids: Layer indices from the target model
66 whose intermediate hidden states are used as auxiliary inputs
67 """
68 pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
69 if config_dict.get("target_hidden_size") is not None:
70 pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"]
71 pre_trained_config["norm_before_residual"] = config_dict.get(
72 "norm_before_residual", False
73 )
74 pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False)
75 peagle_arch_map = {
76 "qwen3": "PeagleQwen3ForCausalLM",
77 "llama": "PeagleLlamaForCausalLM",
78 }
79 model_type = pre_trained_config.get("model_type", "llama")
80 if model_type not in peagle_arch_map:
81 raise ValueError(f"Unsupported model_type {model_type} for PEagle speculator")
82 pre_trained_config["architectures"] = [peagle_arch_map[model_type]]
83 pre_trained_config["pard_token"] = config_dict["mask_token_id"]
84 if config_dict.get("eagle_aux_hidden_state_layer_ids"):
85 pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[
86 "eagle_aux_hidden_state_layer_ids"
87 ]
88
89
90 @register_speculator("dflash")
91 def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
92 """Apply DFlash specific configuration transformations to the `dict` used to
93 construct the Transformers PreTrainedConfig.
94
95 DFlash specific fields:
96 - draft_vocab_size: Size of the draft model's vocabulary
97 - target_hidden_size: Hidden size of the target model
98 - mask_token_id (required): Token ID used for parallel drafting mask
99 placeholders
100 - aux_hidden_state_layer_ids (required): Layer indices from the target
101 model whose intermediate hidden states are used as context for the
102 DFlash drafter. Mapped to both eagle_aux_hidden_state_layer_ids
103 (for gpu_model_runner) and dflash_config.target_layer_ids (for the
104 DFlash model).
105 - sample_from_anchor: Whether to sample from the anchor position. Default
106 False (anchor is a bonus token, only mask tokens predict, yielding
107 block_size - 1 speculative tokens).
108 """
109 pre_trained_config["architectures"] = ["DFlashDraftModel"]
110 pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
111 if config_dict.get("target_hidden_size") is not None:
112 pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"]
113
114 aux_layer_ids = config_dict["aux_hidden_state_layer_ids"]
115 pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids
116
117 # DFlash configs use different indexing for the target layers, see #40727
118 pre_trained_config["dflash_config"] = {
119 "mask_token_id": config_dict["mask_token_id"],
120 "target_layer_ids": [i - 1 for i in aux_layer_ids],
121 "sample_from_anchor": config_dict.get("sample_from_anchor", False),
122 }
123 # Enable causal masking in SWA for vllm-project/speculators models
124 pre_trained_config["dflash_config"]["causal"] = not config_dict.get(
125 "sliding_window_non_causal", True
126 )
127
128
129 @register_speculator("dflash2")
130 def update_dflash2(config_dict: dict, pre_trained_config: dict) -> None:
131 """Apply DFlash2 specific configuration transformations to the `dict` used
132 to construct the Transformers PreTrainedConfig.
133
134 DFlash2 extends DFlash with a grouped convolution layer and a
135 low-rank candidate selector head. It reuses the same DFlash runtime
136 (method="dflash") but has its own architecture (DFlash2DraftModel)
137 and additional config fields for the convolution and selector.
138 """
139 update_dflash(config_dict, pre_trained_config)
140 pre_trained_config["architectures"] = ["DFlash2DraftModel"]
141
142 for key in (
143 "conv_kernel_size",
144 "conv_group_size",
145 "selector_rank",
146 "selector_top_k",
147 ):
148 if config_dict.get(key) is not None:
149 pre_trained_config["dflash_config"][key] = config_dict[key]
150
151
152 @register_speculator("dspark")
153 def update_dspark(config_dict: dict, pre_trained_config: dict) -> None:
154 """Apply DSpark specific configuration transformations to the `dict` used to
155 construct the Transformers PreTrainedConfig.
156
157 DSpark extends DFlash with a Markov logit-bias head, reusing the same
158 Qwen3DSparkModel loader and DSparkSpeculator runtime as the dense DSpark
159 checkpoints (e.g. deepseek-ai/dspark_qwen3_8b_block7).
160
161 DSpark specific fields:
162 - draft_vocab_size: draft vocab size; when smaller than the target vocab the
163 checkpoint also ships d2t/t2d remap tables.
164 - mask_token_id (required): token id for parallel-drafting mask slots.
165 - markov_rank / markov_head_type: low-rank Markov logit-bias head.
166 - block_size: semi-autoregressive draft block size.
167 - enable_confidence_head / confidence_head_with_markov: confidence head.
168 - aux_hidden_state_layer_ids (required): target layer indices feeding the
169 drafter. Mapped to both eagle_aux_hidden_state_layer_ids and
170 target_layer_ids (DSpark's i-1 layer semantics).
171 - sample_from_anchor: Whether to sample from the anchor position. Default
172 False (anchor is a bonus token, only mask tokens predict, yielding
173 block_size - 1 speculative tokens).
174 """
175 architectures = config_dict.get("architectures") or []
176 supported_architectures = {
177 "Qwen3DSparkModel",
178 "Qwen3OmniDSparkModel",
179 }
180 selected_architectures = [
181 architecture
182 for architecture in architectures
183 if architecture in supported_architectures
184 ]
185 # Legacy msModelSpec checkpoints use a training-only architecture name.
186 pre_trained_config["architectures"] = selected_architectures or ["Qwen3DSparkModel"]
187
188 sample_from_anchor = config_dict.get("sample_from_anchor", False)
189 pre_trained_config["sample_from_anchor"] = sample_from_anchor
190 pre_trained_config["dspark_bonus_anchor"] = not sample_from_anchor
191
192 aux_layer_ids = config_dict["aux_hidden_state_layer_ids"]
193 pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids
194 # DSpark indexes target layers as aux_id - 1 (matches the dense configs).
195 pre_trained_config["target_layer_ids"] = [i - 1 for i in aux_layer_ids]
196
197 for key in (
198 "draft_vocab_size",
199 "target_hidden_size",
200 "mask_token_id",
201 "markov_rank",
202 "markov_head_type",
203 "block_size",
204 "enable_confidence_head",
205 "confidence_head_with_markov",
206 "use_aux_hidden_state",
207 ):
208 if config_dict.get(key) is not None:
209 pre_trained_config[key] = config_dict[key]