fix: spec overlap predict shape does not match verify output shapes (#12786)

2025-11-25 19:01:06 -08:00
parent 13e5beeab4
commit f33e5d1ef1
1 changed files with 2 additions and 3 deletions
--- 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
        )