From 4f5ff39bc96da933dd938afc4382793c943ce631 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 21 Jun 2026 15:35:57 -0700 Subject: [PATCH] [Spec] Enable FR-Spec in EAGLE draft-extend CUDA graph by sizing logits buffer from the draft head (#28856) --- .../eagle_draft_extend_cuda_graph_runner.py | 5 +++++ python/sglang/srt/speculative/eagle_worker_v2.py | 16 +++++----------- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index a50d75eed..25aaa58bc 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 463a3fe3a..86af77e7d 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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 - # (speculative_token_map / FR-Spec); fall back to eager in that case. - flashinfer_graph_supported = ( - 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 - ) + supports_cuda_draft_extend_graph = ( + _is_cuda or _is_musa + ) and graph_supported_backend # Capture extend # TODO: support draft extend cuda graph for more attention backends if self.draft_extend_attn_backend and (