Add configurable FlashInfer autotune skips (#31389)

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)

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(