Support loading sample_from_anchor flag from speculators config (#48639)
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Fynn Schmitt-Ulms committed
Jul 20, 2026 at 17:52 UTC
642076d26c98aab899a6cc3dc948856d38c7551b
3 files changed
+24
-6
vllm/transformers_utils/configs/speculators/algos.py
+10
-2
@@ -106,6 +106,9 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
106
DFlash drafter. Mapped to both eagle_aux_hidden_state_layer_ids
107
(for gpu_model_runner) and dflash_config.target_layer_ids (for the
108
DFlash model).
109
+ - sample_from_anchor: Whether to sample from the anchor position. Default
110
+ False (anchor is a bonus token, only mask tokens predict, yielding
111
+ block_size - 1 speculative tokens).
112
"""
113
pre_trained_config["architectures"] = ["DFlashDraftModel"]
114
pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
@@ -119,6 +122,7 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
122
pre_trained_config["dflash_config"] = {
123
"mask_token_id": config_dict["mask_token_id"],
124
"target_layer_ids": [i - 1 for i in aux_layer_ids],
125
+ "sample_from_anchor": config_dict.get("sample_from_anchor", False),
126
}
127
# Enable causal masking in SWA for vllm-project/speculators models
128
pre_trained_config["dflash_config"]["causal"] = not config_dict.get(
@@ -146,10 +150,14 @@ def update_dspark(config_dict: dict, pre_trained_config: dict) -> None:
150
- aux_hidden_state_layer_ids (required): target layer indices feeding the
151
drafter. Mapped to both eagle_aux_hidden_state_layer_ids and
152
target_layer_ids (DSpark's i-1 layer semantics).
153
+ - sample_from_anchor: Whether to sample from the anchor position. Default
154
+ False (anchor is a bonus token, only mask tokens predict, yielding
155
+ block_size - 1 speculative tokens).
156
"""
157
pre_trained_config["architectures"] = ["Qwen3DSparkModel"]
151
- # Speculators DSpark uses the 1+N fill-in block (anchor is a bonus token).
152
- pre_trained_config["dspark_bonus_anchor"] = True
158
+ pre_trained_config["sample_from_anchor"] = config_dict.get(
159
+ "sample_from_anchor", False
160
+ )
161
162
aux_layer_ids = config_dict["aux_hidden_state_layer_ids"]
163
pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+9
@@ -57,6 +57,15 @@ class DFlashSpeculator(DraftModelSpeculator):
57
# Whether the anchor query position is itself a prediction. DFlash default uses
58
# the anchor as the bonus token (only mask tokens predict); DSpark samples from
59
# the anchor and the N-1 mask token positions. See _prepare_dflash_inputs_kernel
60
+ dflash_config = (
61
+ getattr(self.draft_model_config.hf_config, "dflash_config", None) or {}
62
+ )
63
+ if dflash_config.get("sample_from_anchor", False):
64
+ raise ValueError(
65
+ "sample_from_anchor=True is not supported for DFlash. "
66
+ "DFlash uses a fixed 1+N query layout where the anchor "
67
+ "is the bonus token."
68
+ )
69
self.sample_from_anchor = False
70
71
# Context positions for the K/V precompute. Populated by
vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
+5
-4
@@ -40,10 +40,11 @@ class DSparkSpeculator(DFlashSpeculator):
40
def __init__(self, vllm_config: VllmConfig, device: torch.device):
41
super().__init__(vllm_config, device)
42
43
- # Anchor-as-first (N slots) unless the checkpoint uses the 1+N fill-in
44
- # block, where the anchor is a separate bonus token.
45
- self.sample_from_anchor = not getattr(
46
- self.draft_model_config.hf_config, "dspark_bonus_anchor", False
43
+ # Whether to sample from the anchor position. When True, uses anchor-as-first
44
+ # (N slots, each position predicts the next token). When False, uses 1+N
45
+ # fill-in block (anchor is a bonus token).
46
+ self.sample_from_anchor = getattr(
47
+ self.draft_model_config.hf_config, "sample_from_anchor", True
48
)
49
if self.sample_from_anchor:
50
self.num_query_per_req = self.num_speculative_steps