A high-throughput and memory-efficient inference and serving engine for LLMs
[CI/Build][Bugfix] Fix SD LoRA (#45941)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Jun 18, 2026 at 08:34 UTC
ebbb2d55ace74b7066bca0ff8f333012bb8c4299
1 file changed
+8
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+8
@@ -5,6 +5,7 @@ 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.lora.layers.base import BaseLayerWithLoRA
9
from vllm.model_executor.model_loader import get_model
10
11
@@ -48,6 +49,13 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
49
target_embed = getattr(target_inner, "embed_tokens", None) or getattr(
50
target_inner, "embedding", None
51
)
52
+ # If the target's embedding is LoRA-wrapped, share the underlying base
53
+ # layer. The draft is not part of the LoRA adapter; sharing the wrapper
54
+ # would make the draft run the LoRA embedding kernel with the target's
55
+ # punica metadata (sized for the target's token count), causing an
56
+ # out-of-bounds GPU access during multi-step draft decode.
57
+ if isinstance(target_embed, BaseLayerWithLoRA):
58
+ target_embed = target_embed.base_layer
59
draft_embed = getattr(draft_inner, "embed_tokens", None)
60
if target_embed is not None and _should_share(
61
eagle_model, "has_own_embed_tokens", draft_embed, target_embed