diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index 5f27b7fb5..ffef278fb 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -404,9 +404,13 @@ class DFlashDraftModel(nn.Module): pp_proxy_tensors=None, ) -> LogitsProcessorOutput: if input_embeds is None: - raise ValueError( - "DFlashDraftModel requires `input_embeds` (use the target embedding)." - ) + if hasattr(self, "forward_embed"): + input_embeds = self.forward_embed(input_ids) + else: + raise ValueError( + "DFlashDraftModel requires `input_embeds` (use the target " + "embedding)." + ) hidden_states = input_embeds residual: Optional[torch.Tensor] = None diff --git a/python/sglang/srt/models/dspark.py b/python/sglang/srt/models/dspark.py index ea8bcb0ad..0ffc22b3b 100644 --- a/python/sglang/srt/models/dspark.py +++ b/python/sglang/srt/models/dspark.py @@ -376,9 +376,15 @@ class DSparkDraftMixin: def attach_shared_modules( self, *, embed_tokens: nn.Module, lm_head: nn.Module ) -> None: - del embed_tokens + self.embed_tokens = embed_tokens self.lm_head = lm_head + def forward_embed(self, input_ids: torch.Tensor) -> torch.Tensor: + # Embeds with the shared target embedding INSIDE the draft graph + # (the runner skips the eager input_embeds staging when the draft + # model exposes forward_embed). + return self.embed_tokens(input_ids) + def compute_base_logits( self, hidden: torch.Tensor ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: