[Bugfix][Spec Decode] Fix DFlash draft/target layer-count mismatch (#48113)

Signed-off-by: Nick Hill <[email protected]>
Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>

tests/models/registry.py

@@ -1403,7 +1403,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
     # ),
     # [DFlash]
     "DFlashDraftModel": _HfExamplesInfo(
-        "Qwen/Qwen3.5-4B",
+        "Qwen/Qwen3-4B",
         speculative_model="z-lab/Qwen3-4B-DFlash-b16",
         use_original_num_layers=True,  # Need all layers since DFlash has >1 layer,
         max_model_len=8192,  # Reduce max len to ensure test runs in low-VRAM CI env

vllm/model_executor/models/qwen3_dflash.py

@@ -754,6 +754,14 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
         needs_squeeze = hidden_states.dim() == 1
         if needs_squeeze:
             hidden_states = hidden_states.unsqueeze(0)
+        expected = self.model.fc.input_size
+        if hidden_states.shape[-1] != expected:
+            raise ValueError(
+                f"DFlash drafter expects {expected} concatenated aux hidden "
+                f"features but received {hidden_states.shape[-1]}. This usually "
+                "means the draft model's target_layer_ids reference layers that "
+                "do not exist in the target model (incompatible draft/target pair)."
+            )
         result = self.model.fc(hidden_states)
         if needs_squeeze:
             result = result.squeeze(0)