From 435d1c83c1f6cb31503263898342dbecb13a5b16 Mon Sep 17 00:00:00 2001 From: elvischenv <219235043+elvischenv@users.noreply.github.com> Date: Wed, 17 Dec 2025 15:01:39 +0800 Subject: [PATCH] [Perf] Enable Flashinfer autotune by default (#14357) --- docs/advanced_features/server_arguments.md | 1 + python/sglang/srt/layers/moe/fused_moe_triton/layer.py | 3 +++ python/sglang/srt/layers/quantization/fp8.py | 3 +++ .../sglang/srt/layers/quantization/modelopt_quant.py | 1 + python/sglang/srt/model_executor/model_runner.py | 10 +++++++--- python/sglang/srt/server_args.py | 8 ++++---- test/srt/test_deepseek_v3_fp4_4gpu.py | 2 +- 7 files changed, 20 insertions(+), 8 deletions(-) diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index c2c8b1682..587a92d44 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -246,6 +246,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--nsa-prefill` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_decode`, `fa3`, `tilelang`, `aiter` | | `--nsa-decode` | Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding. | `flashmla_kv` | `flashmla_prefill`, `flashmla_kv`, `fa3`, `tilelang`, `aiter` | | `--fp8-gemm-backend` | Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (optimal for Blackwell and low-latency), 'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), 'triton' (fallback, widely compatible), 'aiter' (ROCm only). **NOTE**: This replaces the deprecated environment variables SGLANG_ENABLE_FLASHINFER_FP8_GEMM and SGLANG_SUPPORT_CUTLASS_BLOCK_FP8. | `auto` | `auto`, `deep_gemm`, `flashinfer_trtllm`, `cutlass`, `triton`, `aiter` | +| `--disable-flashinfer-autotune` | Flashinfer autotune is enabled by default. Set this flag to disable the autotune. | `False` | bool flag (set to enable) | ## Speculative decoding | Argument | Description | Defaults | Options | diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 3563d3e3d..b68ec3146 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -54,6 +54,7 @@ from sglang.srt.utils import ( is_cpu, is_flashinfer_available, is_hip, + next_power_of_2, round_up, ) @@ -1087,6 +1088,7 @@ class FlashInferFusedMoE(FusedMoE): local_expert_offset=self.moe_ep_rank * self.num_local_experts, local_num_experts=self.num_local_experts, routing_method_type=self.routing_method_type, + tune_max_num_tokens=next_power_of_2(hidden_states.shape[0]), ) else: @@ -1228,6 +1230,7 @@ class FlashInferFP4MoE(FusedMoE): tile_tokens_dim=None, routing_method_type=routing_method_type, do_finalize=True, + tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]), output=symm_output, )[0] diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index a56c9dc06..830a6752c 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -69,6 +69,7 @@ from sglang.srt.utils import ( is_sm90_supported, is_sm100_supported, log_info_on_rank0, + next_power_of_2, print_warning_once, set_weight_attrs, use_intel_amx_backend, @@ -1384,6 +1385,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): tile_tokens_dim=None, routing_method_type=routing_method_type, use_shuffled_weight=False, + tune_max_num_tokens=next_power_of_2(a_q.shape[0]), ) else: routing_bias_cast = ( @@ -1415,6 +1417,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): ), use_routing_scales_on_input=False, routing_method_type=routing_method_type, + tune_max_num_tokens=next_power_of_2(a_q.shape[0]), ) def maybe_apply_hip_fused_experts( diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 4abb1fda0..c71dd02e2 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -793,6 +793,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): use_routing_scales_on_input=use_routing_scales_on_input, tile_tokens_dim=None, routing_method_type=routing_method_type, + tune_max_num_tokens=next_power_of_2(x.shape[0]), ) from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6f87fdee9..be9833252 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2219,11 +2219,15 @@ class ModelRunner: def _should_run_flashinfer_autotune(self) -> bool: """Check if flashinfer autotune should be run.""" - if not self.server_args.enable_flashinfer_autotune: + if self.server_args.disable_flashinfer_autotune: return False - backend_str = self.server_args.attention_backend - if backend_str not in ["flashinfer", "trtllm_mla", "trtllm_mha"]: + backend_str = self.server_args.moe_runner_backend + if backend_str not in [ + "flashinfer_trtllm", + "flashinfer_cutlass", + "flashinfer_mxfp4", + ]: return False major, _ = torch.cuda.get_device_capability() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6df68aaba..764048180 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -418,7 +418,7 @@ class ServerArgs: fp8_gemm_runner_backend: str = "auto" nsa_prefill_backend: str = "flashmla_sparse" nsa_decode_backend: str = "fa3" - enable_flashinfer_autotune: bool = False + disable_flashinfer_autotune: bool = False # Speculative decoding speculative_algorithm: Optional[str] = None @@ -3313,10 +3313,10 @@ class ServerArgs: "SGLANG_ENABLE_FLASHINFER_FP8_GEMM and SGLANG_SUPPORT_CUTLASS_BLOCK_FP8.", ) parser.add_argument( - "--enable-flashinfer-autotune", - default=ServerArgs.enable_flashinfer_autotune, + "--disable-flashinfer-autotune", + default=ServerArgs.disable_flashinfer_autotune, action="store_true", - help="Enable FlashInfer autotuning for optimal kernel selection.", + help="Disable FlashInfer autotuning.", ) # Speculative decoding diff --git a/test/srt/test_deepseek_v3_fp4_4gpu.py b/test/srt/test_deepseek_v3_fp4_4gpu.py index 2b4f94889..37caaa2f3 100644 --- a/test/srt/test_deepseek_v3_fp4_4gpu.py +++ b/test/srt/test_deepseek_v3_fp4_4gpu.py @@ -16,7 +16,7 @@ from sglang.test.test_utils import ( ) FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3-0324-FP4" -SERVER_LAUNCH_TIMEOUT = 1000 +SERVER_LAUNCH_TIMEOUT = 1200 class TestDeepseekV3FP4(CustomTestCase):