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

Co-authored-by: yukavio <kavioyu@gmail.com>

python/sglang/srt/managers/schedule_batch.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 # Copyright 2023-2024 SGLang Team
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -29,7 +31,7 @@ ScheduleBatch -> ModelWorkerBatch -> ForwardBatch
 
 import dataclasses
 import logging
-from typing import List, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, List, Optional, Set, Tuple, Union
 
 import numpy as np
 import torch
@@ -47,6 +49,10 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import ServerArgs
 
+if TYPE_CHECKING:
+    from sglang.srt.speculative.spec_info import SpecInfo, SpeculativeAlgorithm
+
+
 INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
 
 # Put some global args for easy access
@@ -565,9 +571,13 @@ class ScheduleBatch:
     # Has grammar
     has_grammar: bool = False
 
-    # device
+    # Device
     device: str = "cuda"
 
+    # Speculative decoding
+    spec_info: Optional[SpecInfo] = None
+    spec_algorithm: Optional[SpeculativeAlgorithm] = None
+
     @classmethod
     def init_new(
         cls,
@@ -577,6 +587,7 @@ class ScheduleBatch:
         tree_cache: BasePrefixCache,
         model_config: ModelConfig,
         enable_overlap: bool,
+        speculative_algorithm: Optional[SpeculativeAlgorithm] = None,
     ):
         return cls(
             reqs=reqs,
@@ -589,6 +600,7 @@ class ScheduleBatch:
             has_stream=any(req.stream for req in reqs),
             has_grammar=any(req.grammar for req in reqs),
             device=req_to_token_pool.device,
+            spec_algorithm=speculative_algorithm,
         )
 
     def batch_size(self):
@@ -1103,6 +1115,9 @@ class ScheduleBatch:
         self.has_stream |= other.has_stream
         self.has_grammar |= other.has_grammar
 
+        if self.spec_info:
+            self.spec_info.merge_batch(other.spec_info)
+
     def get_model_worker_batch(self):
         if self.forward_mode.is_decode() or self.forward_mode.is_idle():
             extend_seq_lens = extend_prefix_lens = extend_logprob_start_lens = None
@@ -1144,6 +1159,8 @@ class ScheduleBatch:
             lora_paths=[req.lora_path for req in self.reqs],
             sampling_info=self.sampling_info,
             input_embeds=self.input_embeds,
+            spec_algorithm=self.spec_algorithm,
+            spec_info=self.spec_info,
         )
 
     def copy(self):
@@ -1214,6 +1231,10 @@ class ModelWorkerBatch:
     # The input Embeds
     input_embeds: Optional[torch.tensor] = None
 
+    # Speculative decoding
+    spec_info: Optional[SpecInfo] = None
+    spec_algorithm: Optional[SpeculativeAlgorithm] = None
+
 
 @triton.jit
 def write_req_to_token_pool_triton(