[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

@@ -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
+        # Do NOT cast predict dtype here. The caller (e.g., _draft_extend_for_decode)
-        # gather requires input_ids to have a consistent integer dtype across ranks.
+        # may run this under a plan stream; casting inside the plan stream creates a
-        batch.input_ids = predict.to(torch.int64)
+        # 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,

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,