From 8f0d320d3162f3586c323f2c024b45d0a4fd3fc6 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Thu, 2 Jul 2026 04:36:07 +0800 Subject: [PATCH] [Spec] Enable FlashInfer autotune for spec draft (#29595) --- .../srt/model_executor/runner/base_runner.py | 164 +------------ .../runner/decode_cuda_graph_runner.py | 20 +- .../runner/flashinfer_autotune.py | 224 ++++++++++++++++++ .../eagle_draft_cuda_graph_runner.py | 16 +- .../eagle_draft_extend_cuda_graph_runner.py | 20 +- .../frozen_kv_mtp_cuda_graph_runner.py | 16 +- ...er_eagle_draft_extend_cuda_graph_runner.py | 16 +- .../attention_methods/dense_attention.py | 1 + .../attention_methods/dsa_attention.py | 1 + .../attention_methods/dsv4_attention.py | 1 + .../attention_methods/dual_chunk_attention.py | 1 + .../attention_methods/gdn_attention.py | 1 + .../attention_methods/kda_attention.py | 1 + .../attention_methods/lightning_attention.py | 1 + .../attention_methods/mamba2_attention.py | 1 + .../attention_methods/mla_attention.py | 1 + 16 files changed, 313 insertions(+), 172 deletions(-) create mode 100644 python/sglang/srt/model_executor/runner/flashinfer_autotune.py diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 1be6c95b7..12417c16d 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/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(): + self._dummy_run(batch_size=batch_size, buffers=buffers) - 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: - 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" - ) + run_flashinfer_autotune_forward(self.model_runner, forward_fn, skip_logits=True) 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 - use this -- it reuses an existing runner's buffers via _autotune_buffers - (the eager input registry, or the decode cuda-graph runner's captured - buffers). + sweep; _dummy_run slices it per shape). Eager FlashInfer autotune also + allocates decode-shaped scratch buffers here. Decode cuda-graph autotune + reuses the captured runner buffers instead. """ mr = self.model_runner return _allocate_decode_buffers( diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 3b8e753a9..4c0b3affc 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/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) + 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=getattr( - self.model_runner.attn_backend, - "on_after_cuda_graph_warmup", - None, - ), + post_warmup_hook=post_warmup_hook, ) def recapture_if_needed(self, forward_batch: ForwardBatch): diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py new file mode 100644 index 000000000..5cdc3679e --- /dev/null +++ b/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) diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 3c0ff9c54..480930242 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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( - self.draft_attn_backend, "on_after_cuda_graph_warmup", None - ), + post_warmup_hook=post_warmup_hook, ) def _postprocess_output_to_raw_bs(self, out, raw_bs): 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 d546b0cd4..57b029ba4 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 @@ -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) + 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=getattr( - self.draft_extend_attn_backend, - "on_after_cuda_graph_warmup", - None, - ), + post_warmup_hook=post_warmup_hook, ) def execute(self, forward_batch: ForwardBatch): diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index 785dc6efd..58ef342cb 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/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( - self.draft_attn_backend, "on_after_cuda_graph_warmup", None - ), + post_warmup_hook=post_warmup_hook, ) finally: self.draft_attn_backend.token_to_kv_pool = saved_backend_pool diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 0f8f047d7..5fc143e7d 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/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( - self.attn_backend, "on_after_cuda_graph_warmup", None - ), + post_warmup_hook=post_warmup_hook, ) def replay(self, bs: int, seq_lens_sum: int, spec_info: EagleDraftExtendInput): diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index fb51ed4f5..447811ebf 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index f31a3d328..bcad2b111 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index 9ab3aaa41..282fc847c 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py index 55290bf9b..355e3f30d 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py index f7e8e5b42..fb6bf4606 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py index 70328bd2c..9213e5f2f 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py index 72910077d..43b72fb4b 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py index 998acc31f..09fadaacf 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py index f3db98e5b..641a2d81c 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py +++ b/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