A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 17:45 UTC
[ModelRunner V2] Share identical MTP weights (#42538)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
May 13, 2026 at 11:57 UTC
a505cf807e6d3cb1a3cb71435964e031e95de5fc
1 file changed
+54
-36
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+54
-36
@@ -1,11 +1,29 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
import torch.nn as nn
5
6
from vllm.config import VllmConfig
7
+from vllm.distributed.parallel_state import get_pp_group
8
from vllm.model_executor.model_loader import get_model
9
10
11
+def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
12
+ """Share when the draft has no own copy, or its copy matches the target."""
13
+
14
+ if not getattr(eagle, flag, False) or draft is None:
15
+ return True
16
+ if target is None:
17
+ return False
18
+ # torch.equal on GPU allocates a bool mask the size of the input.
19
+ # Use the faster GPU path when there is plenty of headroom;
20
+ # otherwise compare on CPU.
21
+ w = draft.weight
22
+ if w.is_cuda and torch.cuda.mem_get_info(w.device)[0] < w.numel() * 2:
23
+ return torch.equal(w.cpu(), target.weight.cpu())
24
+ return torch.equal(w, target.weight)
25
+
26
+
27
def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
28
from vllm.compilation.backends import set_model_tag
29
@@ -17,51 +35,51 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
35
vllm_config=vllm_config, model_config=draft_model_config
36
)
37
20
- # Share target embeddings when the draft checkpoint does not include
21
- # its own vocab embedding table.
22
- share_embeddings = True
23
- if hasattr(eagle_model, "has_own_embed_tokens"):
24
- share_embeddings = not eagle_model.has_own_embed_tokens
25
- if share_embeddings:
26
- target_language_model = (
27
- target_model.get_language_model()
28
- if hasattr(target_model, "get_language_model")
29
- else target_model
38
+ target_language_model = (
39
+ target_model.get_language_model()
40
+ if hasattr(target_model, "get_language_model")
41
+ else target_model
42
+ )
43
+ target_inner = target_language_model.model
44
+ draft_inner = eagle_model.model
45
+
46
+ # Skip embedding sharing under PP — each rank owns its own embedding.
47
+ if get_pp_group().world_size == 1:
48
+ target_embed = getattr(target_inner, "embed_tokens", None) or getattr(
49
+ target_inner, "embedding", None
50
)
31
- inner_model = getattr(target_language_model, "model", None)
32
- target_embed_tokens = None
33
- if inner_model is not None:
34
- if hasattr(inner_model, "embed_tokens"):
35
- target_embed_tokens = inner_model.embed_tokens
36
- elif hasattr(inner_model, "embedding"):
37
- target_embed_tokens = inner_model.embedding
38
- if target_embed_tokens is not None and hasattr(eagle_model, "model"):
39
- if hasattr(eagle_model.model, "embed_tokens"):
40
- del eagle_model.model.embed_tokens
41
- eagle_model.model.embed_tokens = target_embed_tokens
51
+ draft_embed = getattr(draft_inner, "embed_tokens", None)
52
+ if target_embed is not None and _should_share(
53
+ eagle_model, "has_own_embed_tokens", draft_embed, target_embed
54
+ ):
55
+ if draft_embed is not None:
56
+ del draft_inner.embed_tokens
57
+ draft_inner.embed_tokens = target_embed
58
43
- # Only share target lm_head when the draft model does not own one.
44
- share_lm_head = True
45
- if hasattr(eagle_model, "has_own_lm_head"):
46
- share_lm_head = not eagle_model.has_own_lm_head
47
- if share_lm_head and hasattr(target_model, "lm_head"):
48
- if hasattr(eagle_model, "lm_head"):
59
+ target_lm_head = getattr(target_model, "lm_head", None)
60
+ draft_lm_head = getattr(eagle_model, "lm_head", None)
61
+ if target_lm_head is not None and _should_share(
62
+ eagle_model, "has_own_lm_head", draft_lm_head, target_lm_head
63
+ ):
64
+ if draft_lm_head is not None:
65
del eagle_model.lm_head
50
- eagle_model.lm_head = target_model.lm_head
66
+ eagle_model.lm_head = target_lm_head
67
52
- # MTP models call compute_logits via shared_head.head (a
53
- # ParallelLMHead inside each MTP layer), not self.model.lm_head.
54
- # If the checkpoint omits a copy of the lm_head weights at the
55
- # MTP layer path, shared_head.head stays uninitialised and
56
- # produces zero/NaN logits. Share it explicitly from the target.
57
- inner = getattr(eagle_model, "model", None)
58
- layers = getattr(inner, "layers", None) if inner is not None else None
68
+ # MTP layers route logits through layer.shared_head.head, not
69
+ # eagle_model.lm_head, so the per-layer copies need fixing up too.
70
+ layers = getattr(draft_inner, "layers", None)
71
if layers is not None:
72
items = layers.values() if isinstance(layers, nn.ModuleDict) else layers
73
for layer in items:
74
sh = getattr(layer, "shared_head", None)
75
if sh is not None and hasattr(sh, "head"):
76
del sh.head
65
- sh.head = target_model.lm_head
77
+ sh.head = target_lm_head
78
+
79
+ # MTP also shares a topk_indices_buffer between target and draft.
80
+ if hasattr(target_inner, "topk_indices_buffer"):
81
+ if hasattr(draft_inner, "topk_indices_buffer"):
82
+ del draft_inner.topk_indices_buffer
83
+ draft_inner.topk_indices_buffer = target_inner.topk_indices_buffer
84
85
return eagle_model