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,