[Bugfix][Spec Decode] Take the DFlash draft's RoPE layout from its own config (#54373)
Signed-off-by: SubSir <tiancaizhangdaxian@sjtu.edu.cn> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Zihan Zhang committed
Aug 31, 2026 at 15:13 UTC
d61b6e1878a60d5a2dafab05565790d0c45ff469
3 files changed
+5
-46
vllm/model_executor/models/qwen3_dflash.py
+5
-25
@@ -75,26 +75,6 @@ def dflash_has_any_non_causal(config: Qwen3Config) -> bool:
75
)
76
77
78
-def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None:
79
- """The target's RoPE layout, from its first attention layer.
80
-
81
- A DFlash head must rotate Q/K the way the target it was distilled against
82
- does, and a mismatch is silent — acceptance collapses but nothing errors and
83
- the output stays correct. Draft checkpoints do not carry this, so take it
84
- from the target. None if the target uses no RoPE.
85
- """
86
- language_model = (
87
- target_model.get_language_model()
88
- if hasattr(target_model, "get_language_model")
89
- else target_model
90
- )
91
- for module in language_model.modules():
92
- style = getattr(module, "is_neox_style", None)
93
- if isinstance(style, bool):
94
- return style
95
- return None
96
-
97
-
78
def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int:
79
spec_config = vllm_config.speculative_config
80
config = spec_config.draft_model_config.hf_config
@@ -316,11 +296,11 @@ class DFlashQwen3DecoderLayer(nn.Module):
296
# non-causal) from the draft config.
297
sliding_window, causal = _resolve_layer_attention(config, layer_idx)
298
319
- # RoPE layout, copied off the target at load time by the draft loader
320
- # (see `dflash_target_rope_is_neox_style`). Checkpoints do not carry it:
321
- # a head distilled from an interleaved-RoPE target must rotate the way
322
- # that target does, or every drafted Q/K is wrong and acceptance
323
- # collapses with no error raised.
299
+ # RoPE layout. The rotation applies to the draft's own Q/K, so this is
300
+ # fixed by how the head was distilled, not by the target: a neox-trained
301
+ # head on an interleaved target still needs neox. A mismatch is silent --
302
+ # acceptance collapses and nothing errors -- so a checkpoint that was
303
+ # distilled the other way has to say so here.
304
is_neox_style = getattr(config, "is_neox_style", True)
305
306
self.self_attn = DFlashQwen3Attention(
vllm/v1/spec_decode/dflash.py
-15
@@ -80,21 +80,6 @@ class DFlashProposer(SpecDecodeBaseProposer):
80
self.draft_model_config.hf_config
81
)
82
83
- @override
84
- def load_model(self, target_model: torch.nn.Module) -> None:
85
- from vllm.model_executor.models.qwen3_dflash import (
86
- dflash_target_rope_is_neox_style,
87
- )
88
-
89
- draft_model_config = self.speculative_config.draft_model_config
90
- assert draft_model_config is not None
91
- # The drafter must rotate Q/K the way its target does. Take that from the
92
- # built target before super() constructs the draft.
93
- is_neox_style = dflash_target_rope_is_neox_style(target_model)
94
- if is_neox_style is not None:
95
- draft_model_config.hf_config.is_neox_style = is_neox_style
96
- super().load_model(target_model)
97
-
83
@override
84
def _create_draft_vllm_config(self) -> VllmConfig:
85
base = super()._create_draft_vllm_config()
vllm/v1/worker/gpu/spec_decode/dflash/utils.py
-6
@@ -15,17 +15,11 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
15
from vllm.compilation.backends import set_model_tag
16
from vllm.model_executor.models.qwen3_dflash import (
17
dflash_has_any_non_causal,
18
- dflash_target_rope_is_neox_style,
18
)
19
20
speculative_config = vllm_config.speculative_config
21
assert speculative_config is not None
22
draft_model_config = speculative_config.draft_model_config
24
- # The drafter must rotate Q/K the way its target does. Take that from the
25
- # built target before super() constructs the draft.
26
- is_neox_style = dflash_target_rope_is_neox_style(target_model)
27
- if is_neox_style is not None:
28
- draft_model_config.hf_config.is_neox_style = is_neox_style
23
# Select an attention backend that supports the drafter's attention: mixing
24
# a non-causal layer onto a causal-only backend would fail.
25
draft_vllm_config = replace(