diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index ffd820701..b54b850f5 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/docs_new/docs/advanced_features/server_arguments.mdx @@ -1451,6 +1451,12 @@ Please consult the documentation below and [server_args.py](https://github.com/s Flashinfer autotune is enabled by default. Set this flag to disable the autotune. `False` bool flag (set to enable) + + + `--flashinfer-autotune-skip-ops` + FlashInfer custom-op identifiers to skip during autotuning. See FlashInfer's autotuning documentation. Skipped ops use the heuristic fallback. SGLang temporarily skips mxfp8_gemm by default due to an IMA. + `None` + string `--radix-cache-backend` diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 46f4bba98..69cf50758 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -30,6 +30,15 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +# TODO: Remove after FlashInfer fixes the mxfp8_gemm autotuning IMA. +FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS = frozenset({"mxfp8_gemm"}) + + +def get_flashinfer_autotune_skip_ops(model_runner: ModelRunner) -> set[str]: + skip_ops = set(model_runner.server_args.flashinfer_autotune_skip_ops or ()) + skip_ops.update(FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS) + return skip_ops + def should_run_flashinfer_autotune( model_runner: ModelRunner, *, for_speculative_draft: bool = False @@ -124,6 +133,9 @@ def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path: str(mr.ps.moe_ep_size), str(mr.model_config.hf_config.__class__.__name__), ] + # A different skip policy must not reuse previously tuned tactics. + skip_ops = get_flashinfer_autotune_skip_ops(mr) + model_key_parts.append("skip_ops=" + ",".join(sorted(skip_ops))) if mr.is_draft_worker: model_key_parts.append(f"draft_quant={mr.model_config.quantization}") model_key = "|".join(model_key_parts) @@ -172,11 +184,11 @@ def flashinfer_autotune_context(model_runner: ModelRunner, *, skip_logits: bool) from sglang.srt.layers.logits_processor import autotune_dummy_run_mode maybe_skip_logits = autotune_dummy_run_mode() + skip_ops = get_flashinfer_autotune_skip_ops(mr) with torch.inference_mode(), autotune( - # Autotuning mxfp8_gemm hits an IMA; skip it. True, cache=str(autotune_cache), - skip_ops={"mxfp8_gemm"}, + skip_ops=skip_ops, ), maybe_skip_logits: yield torch.cuda.current_stream().wait_stream(mr.forward_stream) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 7866b590a..2adfb873f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1739,6 +1739,18 @@ class ServerArgs: disable_flashinfer_autotune: A[ bool, "Disable FlashInfer autotuning.", NS("exec.kernel") ] = False + flashinfer_autotune_skip_ops: A[ + Optional[List[str]], + Arg( + help=( + "FlashInfer custom-op identifiers to skip during autotuning. " + "Skipped ops use FlashInfer's heuristic fallback. SGLang " + "temporarily skips mxfp8_gemm by default due to an IMA." + ), + nargs="+", + ), + NS("exec.kernel"), + ] = None mamba_backend: A[ str, Arg(