diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py index f4b0f73f62..e034cec974 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -120,6 +120,10 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None: "mask_token_id": config_dict["mask_token_id"], "target_layer_ids": [i - 1 for i in aux_layer_ids], } + # Enable causal masking in SWA for vllm-project/speculators models + pre_trained_config["dflash_config"]["causal"] = not config_dict.get( + "sliding_window_non_causal", True + ) @register_speculator("dspark")