From 5351800700e719847951632d7e41776bd71c71b3 Mon Sep 17 00:00:00 2001
From: iridiumine <42236072+iridiumine@users.noreply.github.com>
Date: Mon, 22 Jun 2026 01:26:58 +0800
Subject: [PATCH] [Bugfix] Fix MTP acceptance regression on plan stream by
 moving int64 cast before plan stream context (#28410)

---
 python/sglang/srt/speculative/base_spec_worker.py | 8 +++++---
 python/sglang/srt/speculative/eagle_worker_v2.py  | 6 +++++-
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py
index 9936606d6..acaf61332 100644
--- a/python/sglang/srt/speculative/base_spec_worker.py
+++ b/python/sglang/srt/speculative/base_spec_worker.py
@@ -111,9 +111,11 @@ class EagleDraftWorkerBase(ABC):
         gpu_only = batch.seq_lens_cpu is None
 
         batch.spec_info = draft_extend_input
-        # Normalize draft token ids before ForwardBatch construction; DeepSeekV4 DP
-        # gather requires input_ids to have a consistent integer dtype across ranks.
-        batch.input_ids = predict.to(torch.int64)
+        # Do NOT cast predict dtype here. The caller (e.g., _draft_extend_for_decode)
+        # may run this under a plan stream; casting inside the plan stream creates a
+        # cross-stream dependency that can lead to data races and break MTP acceptance.
+        # The caller should cast to int64 before entering the plan stream context.
+        batch.input_ids = predict
         maybe_detect_oob(
             batch.input_ids,
             0,
diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py
index dfc059d8a..a97490960 100644
--- a/python/sglang/srt/speculative/eagle_worker_v2.py
+++ b/python/sglang/srt/speculative/eagle_worker_v2.py
@@ -799,12 +799,16 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             - 1
         )
 
+        # Cast to int64 before entering plan stream to avoid cross-stream
+        # synchronization issues with .to() inside the plan stream context.
+        next_token_ids = batch_result.next_token_ids.to(torch.int64)
+
         # Prepare for draft extend in a separate stream
         with self.plan_stream_ctx:
             forward_batch = self.prepare_for_draft_extend(
                 draft_extend_input,
                 batch,
-                batch_result.next_token_ids,
+                next_token_ids,
                 self.speculative_num_draft_tokens,
                 self.draft_runner,
                 self.cuda_graph_runner_for_draft_extend,