[Spec] Enable FR-Spec in EAGLE draft-extend CUDA graph by sizing logits buffer from the draft head (#28856)

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -195,6 +195,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
                 global_num_tokens_gpu = None
                 global_num_tokens_for_logprob_gpu = None
 
+            hot_token_id = getattr(self.eagle_worker, "hot_token_id", None)
             if hasattr(
                 self.model_runner.model_config.hf_config, "draft_vocab_size"
             ):  # llama_eagle
@@ -203,6 +204,10 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
                 self.model_runner.model_config.hf_config, "hot_vocab_size"
             ):  # llama_eagle3
                 vocab_size = self.model_runner.model_config.hf_config.hot_vocab_size
+            elif hot_token_id is not None:
+                # FR-Spec: reduced vocab is injected via a late
+                # json_model_override_args, so hf_config lacks it; size from the head.
+                vocab_size = len(hot_token_id)
             else:
                 vocab_size = self.model_runner.model_config.vocab_size
 

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -408,25 +408,19 @@ class EagleDraftWorker(EagleDraftWorkerBase):
                 self.draft_attn_backend, AiterMultiStepDraftBackend
             )
 
-        draft_extend_backend = self.draft_extend_attn_backend
         graph_supported_backend = isinstance(
-            draft_extend_backend,
+            self.draft_extend_attn_backend,
             (
                 TritonAttnBackend,
                 TRTLLMMLABackend,
                 TRTLLMHAAttnBackend,
                 TokenspeedMLABackend,
+                FlashInferAttnBackend,
             ),
         )
-        # FlashInfer draft-extend graph does not support a reduced draft vocab
+        supports_cuda_draft_extend_graph = (
-        # (speculative_token_map / FR-Spec); fall back to eager in that case.
+            _is_cuda or _is_musa
-        flashinfer_graph_supported = (
+        ) and graph_supported_backend
-            isinstance(draft_extend_backend, FlashInferAttnBackend)
-            and self.server_args.speculative_token_map is None
-        )
-        supports_cuda_draft_extend_graph = (_is_cuda or _is_musa) and (
-            graph_supported_backend or flashinfer_graph_supported
-        )
         # Capture extend
         # TODO: support draft extend cuda graph for more attention backends
         if self.draft_extend_attn_backend and (