[Bugfix][Spec Decode] Fix DFlash draft/target layer-count mismatch (#48113)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Nick Hill committed
Jul 10, 2026 at 12:42 UTC
7614b88ebdd9739cd126daf6f33e867e27ce52a7
2 files changed
+9
-1
tests/models/registry.py
+1
-1
@@ -1403,7 +1403,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1403
# ),
1404
# [DFlash]
1405
"DFlashDraftModel": _HfExamplesInfo(
1406
- "Qwen/Qwen3.5-4B",
1406
+ "Qwen/Qwen3-4B",
1407
speculative_model="z-lab/Qwen3-4B-DFlash-b16",
1408
use_original_num_layers=True, # Need all layers since DFlash has >1 layer,
1409
max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env
vllm/model_executor/models/qwen3_dflash.py
+8
@@ -754,6 +754,14 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
754
needs_squeeze = hidden_states.dim() == 1
755
if needs_squeeze:
756
hidden_states = hidden_states.unsqueeze(0)
757
+ expected = self.model.fc.input_size
758
+ if hidden_states.shape[-1] != expected:
759
+ raise ValueError(
760
+ f"DFlash drafter expects {expected} concatenated aux hidden "
761
+ f"features but received {hidden_states.shape[-1]}. This usually "
762
+ "means the draft model's target_layer_ids reference layers that "
763
+ "do not exist in the target model (incompatible draft/target pair)."
764
+ )
765
result = self.model.fc(hidden_states)
766
if needs_squeeze:
767
result = result.squeeze(0)