Dflash integration

vllm/config/speculative.py

@@ -46,7 +46,9 @@
     "pangu_ultra_moe_mtp",
     "step3p5_mtp",
 ]
-EagleModelTypes = Literal["eagle", "eagle3", "extract_hidden_states", MTPModelTypes]
+DFlashModelTypes = Literal["dflash"]
+EagleModelTypes = Literal["eagle", "eagle3", MTPModelTypes, DFlashModelTypes]
+
 NgramGPUTypes = Literal["ngram_gpu"]
 SpeculativeMethod = Literal[
     "ngram",
@@ -196,7 +198,11 @@ def compute_hash(self) -> str:
         factors: list[Any] = []
         # Eagle3 and extract_hidden_states affect the computation graph because
         # they return intermediate hidden states in addition to the final hidden state.
-        uses_aux_hidden_states = self.method in ("eagle3", "extract_hidden_states")
+        uses_aux_hidden_states = self.method in (
+            "eagle3",
+            "extract_hidden_states",
+            "dflash",
+        )
         factors.append(uses_aux_hidden_states)
 
         # The specific layers used also affect the computation graph
@@ -480,7 +486,7 @@ def __post_init__(self):
                 )
 
                 # Automatically detect the method
-                if self.method in ("eagle", "eagle3"):
+                if self.method in ("eagle", "eagle3", "dflash"):
                     pass
                 # examples:
                 # yuhuili/EAGLE-LLaMA3-Instruct-8B
@@ -490,6 +496,8 @@ def __post_init__(self):
                     self.method = "eagle"
                 elif "eagle3" in self.draft_model_config.model.lower():
                     self.method = "eagle3"
+                elif "dflash" in self.draft_model_config.model.lower():
+                    self.method = "dflash"
                 elif self.draft_model_config.hf_config.model_type == "medusa":
                     self.method = "medusa"
                 elif self.draft_model_config.hf_config.model_type == "mlp_speculator":
@@ -795,7 +803,7 @@ def _verify_args(self) -> Self:
             "kimi_k25",
         ]
         if (
-            self.method in ("eagle3", "extract_hidden_states")
+            self.method in ("eagle3", "extract_hidden_states", "dflash")
             and self.target_model_config
             and not any(
                 supported_model in self.target_model_config.hf_text_config.model_type
@@ -843,7 +851,7 @@ def max_num_new_slots_for_drafting(self) -> int:
         return slots_per_req
 
     def use_eagle(self) -> bool:
-        return self.method in ("eagle", "eagle3", "mtp")
+        return self.method in ("eagle", "eagle3", "mtp", "dflash")
 
     def uses_draft_model(self) -> bool:
         return self.method == "draft_model"

vllm/model_executor/models/qwen3.py

@@ -310,9 +310,25 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
     def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
         self.model.aux_hidden_state_layers = layers
 
-    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
+    def build_target_layer_ids(self, num_target_layers: int, num_draft_layers: int):
+        if num_draft_layers == 1:
+            return [num_target_layers // 2]
+        start = 1
+        end = num_target_layers - 3
+        span = end - start
+        target_layer_ids = [
+            int(round(start + (i * span) / (num_draft_layers - 1)))
+            for i in range(num_draft_layers)
+        ]
+        return target_layer_ids
+
+    def get_eagle3_aux_hidden_state_layers(self, method) -> tuple[int, ...]:
         num_layers = len(self.model.layers)
-        return (2, num_layers // 2, num_layers - 3)
+        if method == "dflash":
+            return_layers = self.build_target_layer_ids(num_layers, 5)
+        else:
+            return_layers = [2, num_layers // 2, num_layers - 3]
+        return tuple(return_layers)
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)