Add support for loading dflash2 model in speculators format (#53797)
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Fynn Schmitt-Ulms committed
Aug 27, 2026 at 18:12 UTC
f79a2f5582cb578d333eef73cf146929332f6414
2 files changed
+26
vllm/transformers_utils/configs/speculators/algos.py
+24
@@ -130,6 +130,30 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
130
)
131
132
133
+@register_speculator("dflash2")
134
+def update_dflash2(config_dict: dict, pre_trained_config: dict) -> None:
135
+ """
136
+ Apply DFlash2 specific configuration transformations to the `dict` used
137
+ to construct the Transformers PreTrainedConfig.
138
+
139
+ DFlash2 extends DFlash with a grouped convolution layer and a
140
+ low-rank candidate selector head. It reuses the same DFlash runtime
141
+ (method="dflash") but has its own architecture (DFlash2DraftModel)
142
+ and additional config fields for the convolution and selector.
143
+ """
144
+ update_dflash(config_dict, pre_trained_config)
145
+ pre_trained_config["architectures"] = ["DFlash2DraftModel"]
146
+
147
+ for key in (
148
+ "conv_kernel_size",
149
+ "conv_group_size",
150
+ "selector_rank",
151
+ "selector_top_k",
152
+ ):
153
+ if config_dict.get(key) is not None:
154
+ pre_trained_config["dflash_config"][key] = config_dict[key]
155
+
156
+
157
@register_speculator("dspark")
158
def update_dspark(config_dict: dict, pre_trained_config: dict) -> None:
159
"""
vllm/transformers_utils/configs/speculators/base.py
+2
@@ -134,4 +134,6 @@ class SpeculatorsConfig(PretrainedConfig):
134
}
135
if result["method"] == "peagle":
136
result.update({"method": "eagle3", "parallel_drafting": True})
137
+ elif result["method"] == "dflash2":
138
+ result["method"] = "dflash"
139
return result