From f33e5d1ef150006278c5eba3329ff0d9d8fddaa6 Mon Sep 17 00:00:00 2001
From: timmy-feng <70349932+timmy-feng@users.noreply.github.com>
Date: Tue, 25 Nov 2025 19:01:06 -0800
Subject: [PATCH] fix: spec overlap `predict` shape does not match verify
 output shapes (#12786)

---
 python/sglang/srt/speculative/eagle_info_v2.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py
index 7607af9ac..fec3b00c3 100644
--- a/python/sglang/srt/speculative/eagle_info_v2.py
+++ b/python/sglang/srt/speculative/eagle_info_v2.py
@@ -277,9 +277,8 @@ class EagleVerifyInputV2Mixin:
         device = batch.input_ids.device
 
         candidates = self.draft_token.reshape(bs, self.draft_token_num)
-        predict = torch.zeros(
-            (bs * (self.spec_steps + 1),), dtype=torch.int32, device=device
-        )
+        predict_shape = list(next_token_logits.shape)[:-1]
+        predict = torch.zeros(predict_shape, dtype=torch.int32, device=device).flatten()
         accept_index = torch.full(
             (bs, self.spec_steps + 1), -1, dtype=torch.int32, device=device
         )