[Spec] Enable FlashInfer autotune for spec draft (#29595)

python/sglang/srt/model_executor/runner/base_runner.py

@@ -15,12 +15,9 @@
 
 from __future__ import annotations
 
-import datetime
-import hashlib
 import inspect
 import logging
 from abc import ABC, abstractmethod
-from pathlib import Path
 from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any, Optional, Tuple
 
@@ -43,6 +40,10 @@ from sglang.srt.model_executor.forward_batch_info import (
     PPProxyTensors,
 )
 from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    run_flashinfer_autotune_forward,
+    should_run_flashinfer_autotune,
+)
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.speculative.spec_info import create_dummy_verify_input
 from sglang.srt.utils import (
@@ -213,7 +214,7 @@ class BaseRunner(ABC):
 
         self._pre_initialize_flashinfer_allreduce_workspace()
 
-        if self._should_run_flashinfer_autotune():
+        if should_run_flashinfer_autotune(self.model_runner):
             buffers, batch_size = self._autotune_buffers()
             assert (
                 buffers is not None
@@ -250,83 +251,6 @@ class BaseRunner(ABC):
             dtype=mr.dtype,
         )
 
-    def _should_run_flashinfer_autotune(self) -> bool:
-        """Check if flashinfer autotune should be run."""
-        mr = self.model_runner
-        if mr.server_args.disable_flashinfer_autotune:
-            return False
-
-        # CuteDSL v1 (cutedsl runner + deepep a2a) bypasses MoeRunner and must not
-        # be autotuned -- its _dummy_run would dispatch more tokens per rank than
-        # SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK, tripping a DeepEP assert.
-        # Read server_args directly to avoid depending on initialize_moe_config()
-        # having already populated the MoE backend globals.
-        if (
-            mr.server_args.moe_runner_backend == "flashinfer_cutedsl"
-            and mr.server_args.moe_a2a_backend == "deepep"
-        ):
-            return False
-
-        backend_str = mr.server_args.moe_runner_backend
-
-        # TODO smor- support other cases for flashinfer autotune, such as, mamba backend
-
-        moe_needs_autotune = backend_str in [
-            "flashinfer_trtllm",
-            "flashinfer_trtllm_routed",
-            "flashinfer_mxfp4",
-            "flashinfer_cutedsl",
-            "flashinfer_cutlass",
-        ]
-
-        from sglang.srt.layers.quantization.fp4_utils import (
-            get_fp4_gemm_runner_backend,
-        )
-
-        model_uses_fp4 = mr.model_config.quantization in (
-            "modelopt_fp4",
-            "modelopt_mixed",
-        )
-        fp4_gemm_needs_autotune = model_uses_fp4 and (
-            get_fp4_gemm_runner_backend().is_flashinfer_cutlass()
-            or get_fp4_gemm_runner_backend().is_flashinfer_cutedsl()
-        )
-
-        from sglang.srt.layers.quantization.fp8_utils import (
-            get_fp8_gemm_runner_backend,
-        )
-        from sglang.srt.utils import is_sm100_supported
-
-        model_uses_modelopt_fp8 = mr.model_config.quantization in (
-            "modelopt",
-            "modelopt_fp8",
-            "modelopt_mixed",
-        )
-        # Online MXFP8 (microscaling) linears dispatch to flashinfer's
-        # ``mm_mxfp8``, which the flashinfer fp8 autotune dummy run does not
-        # exercise correctly -- it triggers an illegal memory access inside the
-        # mxfp8 cutlass cubin. The mxfp8 gemm is fixed-config and needs no
-        # tuning, so skip autotune for these models.
-        model_uses_mxfp8 = "mxfp8" in (mr.model_config.quantization or "")
-        fp8_gemm_needs_autotune = not model_uses_mxfp8 and (
-            get_fp8_gemm_runner_backend().is_flashinfer_cutlass()
-            or (model_uses_modelopt_fp8 and is_sm100_supported())
-        )
-
-        if not (
-            moe_needs_autotune or fp4_gemm_needs_autotune or fp8_gemm_needs_autotune
-        ):
-            return False
-
-        major, _ = torch.cuda.get_device_capability()
-        if major < 9:
-            return False
-
-        if mr.spec_algorithm.is_speculative():
-            return not mr.is_draft_worker
-
-        return True
-
     def _flashinfer_autotune(self, *, buffers, batch_size):
         """Run flashinfer autotune.
 
@@ -335,76 +259,11 @@ class BaseRunner(ABC):
         Supplied by warmup() (the decode runner's captured buffers when a graph
         runner exists; a freshly-allocated dummy set in the eager path).
         """
-        from flashinfer.autotuner import autotune
 
-        from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
+        def forward_fn():
-
-        mr = self.model_runner
-        cache_path = self._flashinfer_autotune_cache_path()
-        if envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get():
-            autotune_cache = cache_path
-            logger.info("Running FlashInfer autotune with cache: %s", autotune_cache)
-        else:
-            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-            runs_dir = cache_path.parent / "runs"
-            runs_dir.mkdir(parents=True, exist_ok=True)
-            autotune_cache = (
-                runs_dir / f"{cache_path.stem}.{timestamp}{cache_path.suffix}"
-            )
-            logger.info(
-                "Running FlashInfer autotune (cache reuse DISABLED via "
-                "SGLANG_FLASHINFER_AUTOTUNE_CACHE=0); writing fresh result to: %s",
-                autotune_cache,
-            )
-
-        # Run warmup on the non-default stream to avoid NCCL 2.29+ cudaMemcpyBatchAsync
-        # calls on default stream (unsupported by CUDA) when --enable-symm-mem is used.
-        mr.forward_stream.wait_stream(torch.cuda.current_stream())
-        with torch.get_device_module(mr.device).stream(mr.forward_stream):
-            with (
-                torch.inference_mode(),
-                autotune(True, cache=str(autotune_cache)),
-                autotune_dummy_run_mode(),
-            ):
             self._dummy_run(batch_size=batch_size, buffers=buffers)
-        torch.cuda.current_stream().wait_stream(mr.forward_stream)
-        logger.info("FlashInfer autotune completed.")
 
-    def _flashinfer_autotune_cache_path(self) -> Path:
+        run_flashinfer_autotune_forward(self.model_runner, forward_fn, skip_logits=True)
-        import flashinfer
-
-        mr = self.model_runner
-        major, minor = torch.cuda.get_device_capability(mr.device)
-        arch = f"sm{major}{minor}"
-        flashinfer_version = getattr(flashinfer, "__version__", "unknown")
-
-        server_args = mr.server_args
-        model_key = "|".join(
-            [
-                str(server_args.model_path),
-                str(mr.dtype),
-                str(server_args.quantization),
-                str(server_args.moe_runner_backend),
-                str(mr.tp_size),
-                str(mr.pp_size),
-                str(mr.dp_size),
-                str(mr.moe_ep_size),
-                str(mr.model_config.hf_config.__class__.__name__),
-            ]
-        )
-        cache_key = hashlib.sha256(model_key.encode()).hexdigest()[:16]
-        cache_dir = (
-            Path(envs.SGLANG_CACHE_DIR.get())
-            / "flashinfer"
-            / "autotune"
-            / flashinfer_version
-            / arch
-            / cache_key
-        )
-        cache_dir.mkdir(parents=True, exist_ok=True)
-        return (
-            cache_dir / f"rank_tp{mr.tp_rank}_pp{mr.pp_rank}_dp{mr.dp_rank or 0}.json"
-        )
 
     def _alloc_dummy_decode_buffers(self, max_bs: int, *, num_tokens_per_bs: int = 1):
         """Allocate one static decode-buffer set for a dummy forward, sized to
@@ -413,10 +272,9 @@ class BaseRunner(ABC):
         The PP-parallel DeepGEMM warmup sweeps batch sizes far larger than any
         runner's max_bs (up to ~n_sms*block_m), so no pre-allocated runner buffer
         set fits; it builds one here and hands it to _dummy_run (reused across the
-        sweep; _dummy_run slices it per shape). The flashinfer autotune does NOT
+        sweep; _dummy_run slices it per shape). Eager FlashInfer autotune also
-        use this -- it reuses an existing runner's buffers via _autotune_buffers
+        allocates decode-shaped scratch buffers here. Decode cuda-graph autotune
-        (the eager input registry, or the decode cuda-graph runner's captured
+        reuses the captured runner buffers instead.
-        buffers).
         """
         mr = self.model_runner
         return _allocate_decode_buffers(

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -71,6 +71,9 @@ from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
     freeze_gc,
     get_batch_sizes_to_capture,
 )
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    maybe_flashinfer_autotune_speculative_draft,
+)
 from sglang.srt.model_executor.runner.shape_key import ShapeKey
 from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
     BreakableCudaGraphBackend,
@@ -853,15 +856,22 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
 
             with canary_ctx:
                 shape_key = self._make_graph_key(bs, stream_idx, variant_label)
-                self.backend.capture_one(
-                    shape_key,
-                    run_once,
-                    dummies=None,
                 post_warmup_hook = getattr(
                     self.model_runner.attn_backend,
                     "on_after_cuda_graph_warmup",
                     None,
-                    ),
+                )
+                maybe_flashinfer_autotune_speculative_draft(
+                    self,
+                    run_once,
+                    post_warmup_hook=post_warmup_hook,
+                    skip_logits=False,
+                )
+                self.backend.capture_one(
+                    shape_key,
+                    run_once,
+                    dummies=None,
+                    post_warmup_hook=post_warmup_hook,
                 )
 
     def recapture_if_needed(self, forward_batch: ForwardBatch):

python/sglang/srt/model_executor/runner/flashinfer_autotune.py

@@ -0,0 +1,224 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from __future__ import annotations
+
+import contextlib
+import datetime
+import hashlib
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.model_executor.runner.base_runner import BaseRunner
+
+logger = logging.getLogger(__name__)
+
+
+def should_run_flashinfer_autotune(
+    model_runner: ModelRunner, *, for_speculative_draft: bool = False
+) -> bool:
+    """Check if flashinfer autotune should be run."""
+    mr = model_runner
+    if mr.device != "cuda":
+        return False
+    if mr.server_args.disable_flashinfer_autotune:
+        return False
+
+    # CuteDSL v1 (cutedsl runner + deepep a2a) bypasses MoeRunner and must not
+    # be autotuned -- its _dummy_run would dispatch more tokens per rank than
+    # SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK, tripping a DeepEP assert.
+    # Read server_args directly to avoid depending on initialize_moe_config()
+    # having already populated the MoE backend globals.
+    if (
+        mr.server_args.moe_runner_backend == "flashinfer_cutedsl"
+        and mr.server_args.moe_a2a_backend == "deepep"
+    ):
+        return False
+
+    backend_str = mr.server_args.moe_runner_backend
+
+    # TODO smor- support other cases for flashinfer autotune, such as, mamba backend
+
+    moe_needs_autotune = backend_str in [
+        "flashinfer_trtllm",
+        "flashinfer_trtllm_routed",
+        "flashinfer_mxfp4",
+        "flashinfer_cutedsl",
+        "flashinfer_cutlass",
+    ]
+
+    from sglang.srt.layers.quantization.fp4_utils import (
+        get_fp4_gemm_runner_backend,
+    )
+
+    model_quantization = mr.model_config.quantization
+    model_uses_fp4 = model_quantization in (
+        "modelopt_fp4",
+        "modelopt_mixed",
+    )
+    fp4_gemm_needs_autotune = model_uses_fp4 and (
+        get_fp4_gemm_runner_backend().is_flashinfer_cutlass()
+        or get_fp4_gemm_runner_backend().is_flashinfer_cutedsl()
+    )
+
+    from sglang.srt.layers.quantization.fp8_utils import (
+        get_fp8_gemm_runner_backend,
+    )
+    from sglang.srt.utils import is_sm100_supported
+
+    model_uses_modelopt_fp8 = model_quantization in (
+        "modelopt",
+        "modelopt_fp8",
+        "modelopt_mixed",
+    )
+    # Online MXFP8 (microscaling) linears dispatch to flashinfer's
+    # ``mm_mxfp8``, which the flashinfer fp8 autotune dummy run does not
+    # exercise correctly -- it triggers an illegal memory access inside the
+    # mxfp8 cutlass cubin. The mxfp8 gemm is fixed-config and needs no
+    # tuning, so skip autotune for these models.
+    model_uses_mxfp8 = "mxfp8" in (model_quantization or "")
+    fp8_gemm_needs_autotune = not model_uses_mxfp8 and (
+        get_fp8_gemm_runner_backend().is_flashinfer_cutlass()
+        or (model_uses_modelopt_fp8 and is_sm100_supported())
+    )
+
+    if not (moe_needs_autotune or fp4_gemm_needs_autotune or fp8_gemm_needs_autotune):
+        return False
+
+    if torch.cuda.get_device_capability()[0] < 9:
+        return False
+
+    if mr.spec_algorithm.is_speculative():
+        return mr.is_draft_worker if for_speculative_draft else not mr.is_draft_worker
+
+    return True
+
+
+def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path:
+    import flashinfer
+
+    mr = model_runner
+    major, minor = torch.cuda.get_device_capability(mr.device)
+    arch = f"sm{major}{minor}"
+    flashinfer_version = getattr(flashinfer, "__version__", "unknown")
+
+    server_args = mr.server_args
+    model_key_parts = [
+        str(server_args.model_path),
+        str(mr.dtype),
+        str(server_args.quantization),
+        str(server_args.moe_runner_backend),
+        str(mr.tp_size),
+        str(mr.pp_size),
+        str(mr.dp_size),
+        str(mr.moe_ep_size),
+        str(mr.model_config.hf_config.__class__.__name__),
+    ]
+    if mr.is_draft_worker:
+        model_key_parts.append(f"draft_quant={mr.model_config.quantization}")
+    model_key = "|".join(model_key_parts)
+    cache_key = hashlib.sha256(model_key.encode()).hexdigest()[:16]
+    cache_dir = (
+        Path(envs.SGLANG_CACHE_DIR.get())
+        / "flashinfer"
+        / "autotune"
+        / flashinfer_version
+        / arch
+        / cache_key
+    )
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir / f"rank_tp{mr.tp_rank}_pp{mr.pp_rank}_dp{mr.dp_rank or 0}.json"
+
+
+@contextlib.contextmanager
+def flashinfer_autotune_context(model_runner: ModelRunner, *, skip_logits: bool):
+    from flashinfer.autotuner import autotune
+
+    mr = model_runner
+    cache_path = flashinfer_autotune_cache_path(mr)
+    if envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get():
+        autotune_cache = cache_path
+        logger.info("Running FlashInfer autotune with cache: %s", autotune_cache)
+    else:
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        runs_dir = cache_path.parent / "runs"
+        runs_dir.mkdir(parents=True, exist_ok=True)
+        autotune_cache = runs_dir / f"{cache_path.stem}.{timestamp}{cache_path.suffix}"
+        logger.info(
+            "Running FlashInfer autotune (cache reuse DISABLED via "
+            "SGLANG_FLASHINFER_AUTOTUNE_CACHE=0); writing fresh result to: %s",
+            autotune_cache,
+        )
+
+    # Run warmup on the non-default stream to avoid NCCL 2.29+ cudaMemcpyBatchAsync
+    # calls on default stream (unsupported by CUDA) when --enable-symm-mem is used.
+    mr.forward_stream.wait_stream(torch.cuda.current_stream())
+    with torch.get_device_module(mr.device).stream(mr.forward_stream):
+        maybe_skip_logits = contextlib.nullcontext()
+        if skip_logits:
+            from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
+
+            maybe_skip_logits = autotune_dummy_run_mode()
+        with torch.inference_mode(), autotune(
+            True, cache=str(autotune_cache)
+        ), maybe_skip_logits:
+            yield
+    torch.cuda.current_stream().wait_stream(mr.forward_stream)
+    logger.info("FlashInfer autotune completed.")
+
+
+def run_flashinfer_autotune_forward(
+    model_runner: ModelRunner, forward_fn: Callable[[], None], *, skip_logits: bool
+) -> None:
+    """Run flashinfer autotune forward."""
+    with flashinfer_autotune_context(model_runner, skip_logits=skip_logits):
+        forward_fn()
+
+
+def maybe_flashinfer_autotune_speculative_draft(
+    runner: BaseRunner,
+    forward_fn: Callable[[], None],
+    *,
+    post_warmup_hook: Optional[Callable[[], None]] = None,
+    skip_logits: bool = False,
+) -> None:
+    """Run speculative draft flashinfer autotune."""
+    mr = runner.model_runner
+    phase_key = f"{runner.__class__.__module__}.{runner.__class__.__qualname__}"
+    tuned_phases = getattr(mr, "_flashinfer_spec_draft_autotuned_phases", None)
+    if tuned_phases is None:
+        tuned_phases = set()
+        mr._flashinfer_spec_draft_autotuned_phases = tuned_phases
+    if phase_key in tuned_phases:
+        return
+    if (
+        not mr.spec_algorithm.is_speculative()
+        or not mr.is_draft_worker
+        or not should_run_flashinfer_autotune(mr, for_speculative_draft=True)
+    ):
+        return
+
+    def run_and_reset():
+        forward_fn()
+        if post_warmup_hook is not None:
+            post_warmup_hook()
+
+    run_flashinfer_autotune_forward(mr, run_and_reset, skip_logits=skip_logits)
+    tuned_phases.add(phase_key)

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -28,6 +28,9 @@ from sglang.srt.model_executor.runner import (
     get_batch_sizes_to_capture,
     model_capture_mode,
 )
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    maybe_flashinfer_autotune_speculative_draft,
+)
 from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
 from sglang.srt.model_executor.runner_backend_utils import (
     CUDA_GRAPH_CAPTURE_FAILED_MSG,
@@ -440,13 +443,20 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
             forward_batch.mark_forward_metadata_ready()
             self.deepep_adapter.capture(is_extend_in_batch=False)
             shape_key = self._make_graph_key(num_seqs)
+            post_warmup_hook = getattr(
+                self.draft_attn_backend, "on_after_cuda_graph_warmup", None
+            )
+            maybe_flashinfer_autotune_speculative_draft(
+                self,
+                run_once,
+                post_warmup_hook=post_warmup_hook,
+                skip_logits=False,
+            )
             self.backend.capture_one(
                 shape_key,
                 run_once,
                 dummies=None,
-                post_warmup_hook=getattr(
+                post_warmup_hook=post_warmup_hook,
-                    self.draft_attn_backend, "on_after_cuda_graph_warmup", None
-                ),
             )
 
     def _postprocess_output_to_raw_bs(self, out, raw_bs):

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -28,6 +28,9 @@ from sglang.srt.model_executor.runner import (
     get_batch_sizes_to_capture,
     model_capture_mode,
 )
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    maybe_flashinfer_autotune_speculative_draft,
+)
 from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
 from sglang.srt.model_executor.runner_backend_utils import (
     CUDA_GRAPH_CAPTURE_FAILED_MSG,
@@ -432,15 +435,22 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             )
             with canary_ctx:
                 shape_key = self._make_graph_key(bs)
-                self.backend.capture_one(
-                    shape_key,
-                    run_once,
-                    dummies=None,
                 post_warmup_hook = getattr(
                     self.draft_extend_attn_backend,
                     "on_after_cuda_graph_warmup",
                     None,
-                    ),
+                )
+                maybe_flashinfer_autotune_speculative_draft(
+                    self,
+                    run_once,
+                    post_warmup_hook=post_warmup_hook,
+                    skip_logits=False,
+                )
+                self.backend.capture_one(
+                    shape_key,
+                    run_once,
+                    dummies=None,
+                    post_warmup_hook=post_warmup_hook,
                 )
 
     def execute(self, forward_batch: ForwardBatch):

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py

@@ -25,6 +25,9 @@ from sglang.srt.model_executor.runner import (
     get_batch_sizes_to_capture,
     model_capture_mode,
 )
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    maybe_flashinfer_autotune_speculative_draft,
+)
 from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
 from sglang.srt.model_executor.runner_backend_utils import (
     CUDA_GRAPH_CAPTURE_FAILED_MSG,
@@ -321,13 +324,20 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
                 )
                 self.deepep_adapter.capture(is_extend_in_batch=False)
                 shape_key = self._make_graph_key(request_bs)
+                post_warmup_hook = getattr(
+                    self.draft_attn_backend, "on_after_cuda_graph_warmup", None
+                )
+                maybe_flashinfer_autotune_speculative_draft(
+                    self,
+                    run_once,
+                    post_warmup_hook=post_warmup_hook,
+                    skip_logits=False,
+                )
                 self.backend.capture_one(
                     shape_key,
                     run_once,
                     dummies=None,
-                    post_warmup_hook=getattr(
+                    post_warmup_hook=post_warmup_hook,
-                        self.draft_attn_backend, "on_after_cuda_graph_warmup", None
-                    ),
                 )
         finally:
             self.draft_attn_backend.token_to_kv_pool = saved_backend_pool

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -52,6 +52,9 @@ from sglang.srt.model_executor.runner import (
     get_batch_sizes_to_capture,
     model_capture_mode,
 )
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    maybe_flashinfer_autotune_speculative_draft,
+)
 from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
 from sglang.srt.model_executor.runner_backend_utils import (
     CUDA_GRAPH_CAPTURE_FAILED_MSG,
@@ -369,13 +372,20 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True)
             self.deepep_adapter.capture(is_extend_in_batch=True)
             shape_key = self._make_graph_key(bs)
+            post_warmup_hook = getattr(
+                self.attn_backend, "on_after_cuda_graph_warmup", None
+            )
+            maybe_flashinfer_autotune_speculative_draft(
+                self,
+                run_once,
+                post_warmup_hook=post_warmup_hook,
+                skip_logits=False,
+            )
             self.backend.capture_one(
                 shape_key,
                 run_once,
                 dummies=None,
-                post_warmup_hook=getattr(
+                post_warmup_hook=post_warmup_hook,
-                    self.attn_backend, "on_after_cuda_graph_warmup", None
-                ),
             )
 
     def replay(self, bs: int, seq_lens_sum: int, spec_info: EagleDraftExtendInput):

python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py

@@ -276,6 +276,7 @@ class TinyModelConfig:
         self.is_encoder_decoder = False
         self.is_multimodal = False
         self.is_generation = True
+        self.quantization = None
         self.is_hybrid_swa = sliding_window_size is not None
         self.is_local_attention_model = sliding_window_size is not None
         self.attention_chunk_size = None

python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py

@@ -242,6 +242,7 @@ class TinyDSAModelConfig:
         self.is_encoder_decoder = False
         self.is_multimodal = False
         self.is_generation = True
+        self.quantization = None
         self.is_hybrid_swa = False
         self.attention_chunk_size = None
         self.sliding_window_size = None

python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py

@@ -265,6 +265,7 @@ class TinyDSV4ModelConfig:
         self.is_encoder_decoder = False
         self.is_multimodal = False
         self.is_generation = True
+        self.quantization = None
         self.is_hybrid_swa = False
         self.is_local_attention_model = False
         self.attention_chunk_size = None

python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py

@@ -274,6 +274,7 @@ class TinyDualChunkModelConfig:
         self.is_encoder_decoder = False
         self.is_multimodal = False
         self.is_generation = True
+        self.quantization = None
         self.is_hybrid_swa = False
         self.attention_chunk_size = None
         self.sliding_window_size = None

python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py

@@ -178,6 +178,7 @@ class TinyGDNModelConfig:
         self.is_encoder_decoder = False
         self.is_multimodal = False
         self.is_generation = True
+        self.quantization = None
         self.is_hybrid_swa = False
         self.is_local_attention_model = False
         self.attention_chunk_size = None

python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py

@@ -184,6 +184,7 @@ class TinyKDAModelConfig:
         self.is_encoder_decoder = False
         self.is_multimodal = False
         self.is_generation = True
+        self.quantization = None
         self.is_hybrid_swa = False
         self.is_local_attention_model = False
         self.attention_chunk_size = None

python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py

@@ -187,6 +187,7 @@ class TinyLightningModelConfig:
         self.is_encoder_decoder = False
         self.is_multimodal = False
         self.is_generation = True
+        self.quantization = None
         self.is_hybrid_swa = False
         self.is_local_attention_model = False
         self.attention_chunk_size = None

python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py

@@ -275,6 +275,7 @@ class TinyMamba2ModelConfig:
         self.is_encoder_decoder = False
         self.is_multimodal = False
         self.is_generation = True
+        self.quantization = None
         self.is_hybrid_swa = False
         self.is_local_attention_model = False
         self.attention_chunk_size = None

python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py

@@ -178,6 +178,7 @@ class TinyMLAModelConfig:
         self.is_encoder_decoder = False
         self.is_multimodal = False
         self.is_generation = True
+        self.quantization = None
         self.is_hybrid_swa = False
         self.is_local_attention_model = False
         self.attention_chunk_size = None