Add draft extend CUDA graph for flashinfer backend (#6805)

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -80,7 +80,9 @@ class EAGLEDraftExtendCudaGraphRunner:
 
             self.seq_lens = torch.ones((self.max_bs,), dtype=torch.int32)
             self.extend_seq_lens = torch.ones((self.max_bs,), dtype=torch.int32)
-            self.accept_length = torch.ones((self.max_bs,), dtype=torch.int32)
+            self.accept_length = (
+                torch.ones((self.max_bs,), dtype=torch.int32) * self.num_tokens_per_bs
+            )
 
         # Capture
         try:

python/sglang/srt/speculative/eagle_worker.py

@@ -156,6 +156,7 @@ class EAGLEWorker(TpModelWorker):
         if self.server_args.attention_backend == "flashinfer":
             if not global_server_args_dict["use_mla_backend"]:
                 from sglang.srt.layers.attention.flashinfer_backend import (
+                    FlashInferAttnBackend,
                     FlashInferMultiStepDraftBackend,
                 )
 
@@ -164,8 +165,13 @@ class EAGLEWorker(TpModelWorker):
                     self.topk,
                     self.speculative_num_steps,
                 )
+                self.draft_extend_attn_backend = FlashInferAttnBackend(
+                    self.draft_model_runner,
+                    skip_prefill=False,
+                )
             else:
                 from sglang.srt.layers.attention.flashinfer_mla_backend import (
+                    FlashInferMLAAttnBackend,
                     FlashInferMLAMultiStepDraftBackend,
                 )
 
@@ -174,7 +180,10 @@ class EAGLEWorker(TpModelWorker):
                     self.topk,
                     self.speculative_num_steps,
                 )
-            self.draft_extend_attn_backend = None
+                self.draft_extend_attn_backend = FlashInferMLAAttnBackend(
+                    self.draft_model_runner,
+                    skip_prefill=False,
+                )
             self.padded_static_len = self.speculative_num_steps + 1
             self.has_prefill_wrapper_verify = True
         elif self.server_args.attention_backend == "triton":