[Perf] Fold dspark dense draft embedding into the draft graph via forward_embed (#31985)

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
 

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]]: