[Bugfix][DeepSeek-V4] Fix Spec V2 Draft Input ID Dtype for DP Collectives (#28043)

Co-authored-by: zhujunyu <zhujunyu.666@bytedance.com>
2026-06-14 21:11:53 -07:00
parent bf38a0b03d
commit ce9fad7196
1 changed files with 3 additions and 1 deletions
@@ -111,7 +111,9 @@ class EagleDraftWorkerBase(ABC):
        gpu_only = batch.seq_lens_cpu is None

        batch.spec_info = draft_extend_input
-        batch.input_ids = predict
+        # 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)
        maybe_detect_oob(
            batch.input_ids,
            0,