Eagle speculative decoding part 2: Fix cuda graph + DP attention hanging (#2684)

Co-authored-by: yukavio <kavioyu@gmail.com>
2024-12-31 02:25:05 -08:00
parent 6c42fa229d
commit b0524c3789
7 changed files with 131 additions and 58 deletions
@@ -146,7 +146,10 @@ class LogitsProcessor(nn.Module):

        # Compute logits
        last_logits = self._get_logits(last_hidden, lm_head)
-        if not logits_metadata.extend_return_logprob:
+        if (
+            not logits_metadata.extend_return_logprob
+            or logits_metadata.capture_hidden_mode.need_capture()
+        ):
            # Decode mode or extend mode without return_logprob.
            return LogitsProcessorOutput(
                next_token_logits=last_logits,