[Bugfix] DFlash fc sized wrong when num_target_layers != num_hidden_layers (#48524)

Signed-off-by: mgoin <[email protected]>
Co-authored-by: Codex <[email protected]>

tests/v1/spec_decode/test_dflash_causality.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Config-only resolution of DFlash draft attention causality.
+"""Config-only DFlash behavior.
 
 ``dflash_has_any_non_causal`` decides pre-build whether the draft needs a
 non-causal-capable backend, so its branch table (explicit override, SWA-derived
@@ -13,8 +13,12 @@ import pytest
 
 from vllm.model_executor.models.qwen3_dflash import (
     _dflash_layer_causal,
+    _get_dflash_fc_input_size,
     dflash_has_any_non_causal,
 )
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
+    get_eagle3_aux_layers_from_config,
+)
 
 
 def _config(num_hidden_layers, layer_types=None, causal_override=None):
@@ -53,3 +57,35 @@ def test_dflash_layer_causal_is_per_layer():
     config = _config(2, layer_types=["sliding_attention", "full_attention"])
     assert _dflash_layer_causal(config, 0) is True
     assert _dflash_layer_causal(config, 1) is False
+
+
+def _vllm_config(**draft_config):
+    config = SimpleNamespace(**draft_config)
+    return SimpleNamespace(
+        speculative_config=SimpleNamespace(
+            draft_model_config=SimpleNamespace(hf_config=config)
+        )
+    )
+
+
+def test_dflash_fc_uses_aux_layer_count():
+    vllm_config = _vllm_config(
+        num_hidden_layers=5,
+        hidden_size=4096,
+        target_hidden_size=None,
+        target_layer_ids=[1, 17, 32],
+    )
+
+    assert _get_dflash_fc_input_size(vllm_config) == 3 * 4096
+
+
+@pytest.mark.parametrize("config_name", ["dflash_config", "eagle_config"])
+def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name):
+    layer_ids = [1, 17, 32]
+    vllm_config = _vllm_config(
+        **{config_name: {"layer_ids": layer_ids}},
+    )
+
+    assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == tuple(
+        layer_ids
+    )

vllm/model_executor/models/qwen3_dflash.py

@@ -35,6 +35,9 @@ from vllm.multimodal.inputs import NestedTensors
 from vllm.transformers_utils.config import set_default_rope_theta
 from vllm.transformers_utils.repo_utils import get_hf_file_bytes
 from vllm.v1.attention.backend import AttentionType
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
+    get_eagle3_aux_layers_from_config,
+)
 
 from .qwen2 import Qwen2MLP as Qwen3MLP
 from .qwen3 import Qwen3ForCausalLM
@@ -69,6 +72,17 @@ def dflash_has_any_non_causal(config: Qwen3Config) -> bool:
     )
 
 
+def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int:
+    spec_config = vllm_config.speculative_config
+    config = spec_config.draft_model_config.hf_config
+    aux_layers = get_eagle3_aux_layers_from_config(spec_config)
+    num_features_to_use = len(aux_layers) if aux_layers else config.num_hidden_layers
+    target_hidden_size = (
+        getattr(config, "target_hidden_size", None) or config.hidden_size
+    )
+    return target_hidden_size * num_features_to_use
+
+
 def _resolve_layer_attention(
     config: Qwen3Config, layer_idx: int
 ) -> tuple[int | None, bool]:
@@ -395,17 +409,10 @@ class DFlashQwen3Model(nn.Module):
             ]
         )
         if self.use_aux_hidden_state:
-            num_features_to_use = self.config.num_hidden_layers
-            if "target_layer_ids" in drafter_config:
-                num_features_to_use = len(drafter_config["target_layer_ids"])
-            elif "layer_ids" in drafter_config:
-                num_features_to_use = len(drafter_config["layer_ids"])
-            if hasattr(self.config, "target_hidden_size"):
-                fc_input_size = self.config.target_hidden_size * num_features_to_use
-            else:
-                fc_input_size = self.config.hidden_size * num_features_to_use
             self.fc = ReplicatedLinear(
-                input_size=fc_input_size,
+                input_size=_get_dflash_fc_input_size(
+                    vllm_config,
+                ),
                 output_size=self.config.hidden_size,
                 bias=False,
                 params_dtype=vllm_config.model_config.dtype,

vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py

@@ -53,6 +53,13 @@ def get_eagle3_aux_layers_from_config(
         target_layer_ids = getattr(hf_config, "target_layer_ids", None)
         if target_layer_ids:
             layer_ids = [i + 1 for i in target_layer_ids]
+    if not layer_ids:
+        for config_name in ("dflash_config", "eagle_config"):
+            drafter_config = getattr(hf_config, config_name, None)
+            if drafter_config and isinstance(drafter_config, dict):
+                layer_ids = drafter_config.get("layer_ids")
+                if layer_ids:
+                    break
     if layer_ids and isinstance(layer_ids, (list, tuple)):
         return tuple(layer_ids)
     return None