[Model Runer V2][DFlash] Fix lm head sharing for dflash (#46435)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Jun 23, 2026 at 14:09 UTC
ef361de9163e1d0fe8c0ea9028b320994c24f218
1 file changed
+4
-7
vllm/v1/worker/gpu/spec_decode/dflash/utils.py
+4
-7
@@ -55,16 +55,13 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
55
del draft_inner.embed_tokens
56
draft_inner.embed_tokens = target_embed
57
58
- # Share lm_head with the target unless the draft remaps vocab via
59
- # draft_id_to_target_id (in which case its own lm_head is required).
58
target_lm_head = getattr(target_model, "lm_head", None)
59
draft_lm_head = getattr(dflash_model, "lm_head", None)
62
- if (
63
- target_lm_head is not None
64
- and draft_lm_head is not None
65
- and getattr(dflash_model, "draft_id_to_target_id", None) is None
60
+ if target_lm_head is not None and _should_share(
61
+ dflash_model, "has_own_lm_head", draft_lm_head, target_lm_head
62
):
67
- del dflash_model.lm_head
63
+ if draft_lm_head is not None:
64
+ del dflash_model.lm_head
65
dflash_model.lm_head = target_lm_head
66
67
return dflash_model