[Bugfix] DFlash fc sized wrong when num_target_layers != num_hidden_layers (#48524)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Codex <noreply@openai.com>
Michael Goin committed
Jul 21, 2026 at 17:42 UTC
a7d00ec051624e551ea822ec55d1113d117e47b7
3 files changed
+61
-11
tests/v1/spec_decode/test_dflash_causality.py
+37
-1
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Config-only resolution of DFlash draft attention causality.
3
+"""Config-only DFlash behavior.
4
5
``dflash_has_any_non_causal`` decides pre-build whether the draft needs a
6
non-causal-capable backend, so its branch table (explicit override, SWA-derived
@@ -13,8 +13,12 @@ import pytest
13
14
from vllm.model_executor.models.qwen3_dflash import (
15
_dflash_layer_causal,
16
+ _get_dflash_fc_input_size,
17
dflash_has_any_non_causal,
18
)
19
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
20
+ get_eagle3_aux_layers_from_config,
21
+)
22
23
24
def _config(num_hidden_layers, layer_types=None, causal_override=None):
@@ -53,3 +57,35 @@ def test_dflash_layer_causal_is_per_layer():
57
config = _config(2, layer_types=["sliding_attention", "full_attention"])
58
assert _dflash_layer_causal(config, 0) is True
59
assert _dflash_layer_causal(config, 1) is False
60
+
61
+
62
+def _vllm_config(**draft_config):
63
+ config = SimpleNamespace(**draft_config)
64
+ return SimpleNamespace(
65
+ speculative_config=SimpleNamespace(
66
+ draft_model_config=SimpleNamespace(hf_config=config)
67
+ )
68
+ )
69
+
70
+
71
+def test_dflash_fc_uses_aux_layer_count():
72
+ vllm_config = _vllm_config(
73
+ num_hidden_layers=5,
74
+ hidden_size=4096,
75
+ target_hidden_size=None,
76
+ target_layer_ids=[1, 17, 32],
77
+ )
78
+
79
+ assert _get_dflash_fc_input_size(vllm_config) == 3 * 4096
80
+
81
+
82
+@pytest.mark.parametrize("config_name", ["dflash_config", "eagle_config"])
83
+def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name):
84
+ layer_ids = [1, 17, 32]
85
+ vllm_config = _vllm_config(
86
+ **{config_name: {"layer_ids": layer_ids}},
87
+ )
88
+
89
+ assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == tuple(
90
+ layer_ids
91
+ )
vllm/model_executor/models/qwen3_dflash.py
+17
-10
@@ -35,6 +35,9 @@ from vllm.multimodal.inputs import NestedTensors
35
from vllm.transformers_utils.config import set_default_rope_theta
36
from vllm.transformers_utils.repo_utils import get_hf_file_bytes
37
from vllm.v1.attention.backend import AttentionType
38
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
39
+ get_eagle3_aux_layers_from_config,
40
+)
41
42
from .qwen2 import Qwen2MLP as Qwen3MLP
43
from .qwen3 import Qwen3ForCausalLM
@@ -69,6 +72,17 @@ def dflash_has_any_non_causal(config: Qwen3Config) -> bool:
72
)
73
74
75
+def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int:
76
+ spec_config = vllm_config.speculative_config
77
+ config = spec_config.draft_model_config.hf_config
78
+ aux_layers = get_eagle3_aux_layers_from_config(spec_config)
79
+ num_features_to_use = len(aux_layers) if aux_layers else config.num_hidden_layers
80
+ target_hidden_size = (
81
+ getattr(config, "target_hidden_size", None) or config.hidden_size
82
+ )
83
+ return target_hidden_size * num_features_to_use
84
+
85
+
86
def _resolve_layer_attention(
87
config: Qwen3Config, layer_idx: int
88
) -> tuple[int | None, bool]:
@@ -395,17 +409,10 @@ class DFlashQwen3Model(nn.Module):
409
]
410
)
411
if self.use_aux_hidden_state:
398
- num_features_to_use = self.config.num_hidden_layers
399
- if "target_layer_ids" in drafter_config:
400
- num_features_to_use = len(drafter_config["target_layer_ids"])
401
- elif "layer_ids" in drafter_config:
402
- num_features_to_use = len(drafter_config["layer_ids"])
403
- if hasattr(self.config, "target_hidden_size"):
404
- fc_input_size = self.config.target_hidden_size * num_features_to_use
405
- else:
406
- fc_input_size = self.config.hidden_size * num_features_to_use
412
self.fc = ReplicatedLinear(
408
- input_size=fc_input_size,
413
+ input_size=_get_dflash_fc_input_size(
414
+ vllm_config,
415
+ ),
416
output_size=self.config.hidden_size,
417
bias=False,
418
params_dtype=vllm_config.model_config.dtype,
vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
+7
@@ -53,6 +53,13 @@ def get_eagle3_aux_layers_from_config(
53
target_layer_ids = getattr(hf_config, "target_layer_ids", None)
54
if target_layer_ids:
55
layer_ids = [i + 1 for i in target_layer_ids]
56
+ if not layer_ids:
57
+ for config_name in ("dflash_config", "eagle_config"):
58
+ drafter_config = getattr(hf_config, config_name, None)
59
+ if drafter_config and isinstance(drafter_config, dict):
60
+ layer_ids = drafter_config.get("layer_ids")
61
+ if layer_ids:
62
+ break
63
if layer_ids and isinstance(layer_ids, (list, tuple)):
64
return tuple(layer_ids)
65
return None