[Refactor] Add -fp4-gemm-backend to replace SGLANG_FLASHINFER_FP4_GEMM_BACKEND (#16534)

Co-authored-by: Vincent Zhong <207368749+vincentzed@users.noreply.github.com>

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py

@@ -15,8 +15,8 @@ from sglang.srt.layers.parameter import (
 from sglang.srt.layers.quantization.compressed_tensors.schemes import (
     CompressedTensorsScheme,
 )
+from sglang.srt.layers.quantization.fp4_utils import get_fp4_gemm_runner_backend
 from sglang.srt.layers.quantization.modelopt_quant import (
-    FLASHINFER_FP4_GEMM_BACKEND,
     enable_flashinfer_fp4_gemm,
     fp4_gemm,
     fp4_quantize,
@@ -98,7 +98,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
             layer.weight_global_scale.max().to(torch.float32), requires_grad=False
         )
 
-        if FLASHINFER_FP4_GEMM_BACKEND == "trtllm":
+        if get_fp4_gemm_runner_backend().is_trtllm():
             # FlashInfer TRTLLM FP4 GEMM requires a different weight layout.
             # FlashInfer provides nvfp4_quantize to quantize + shuffle the
             # layout but we use our own quantization so we have to call

python/sglang/srt/layers/quantization/fp4_utils.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import logging
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+class Fp4GemmRunnerBackend(Enum):
+    """Enum for FP4 GEMM runner backend selection."""
+
+    AUTO = "auto"
+    CUDNN = "cudnn"
+    CUTLASS = "cutlass"
+    TRTLLM = "trtllm"
+
+    def is_auto(self) -> bool:
+        return self == Fp4GemmRunnerBackend.AUTO
+
+    def is_cudnn(self) -> bool:
+        return self == Fp4GemmRunnerBackend.CUDNN
+
+    def is_cutlass(self) -> bool:
+        return self == Fp4GemmRunnerBackend.CUTLASS
+
+    def is_trtllm(self) -> bool:
+        return self == Fp4GemmRunnerBackend.TRTLLM
+
+
+FP4_GEMM_RUNNER_BACKEND: Fp4GemmRunnerBackend | None = None
+
+
+def initialize_fp4_gemm_config(server_args: ServerArgs) -> None:
+    """Initialize FP4 GEMM configuration from server args."""
+    global FP4_GEMM_RUNNER_BACKEND
+
+    backend = server_args.fp4_gemm_runner_backend
+
+    # Handle deprecated env var for backward compatibility
+    # TODO: Remove this in a future version
+    if envs.SGLANG_FLASHINFER_FP4_GEMM_BACKEND.is_set():
+        env_backend = envs.SGLANG_FLASHINFER_FP4_GEMM_BACKEND.get()
+        if backend == "auto":
+            logger.warning(
+                "SGLANG_FLASHINFER_FP4_GEMM_BACKEND is deprecated. "
+                f"Please use '--fp4-gemm-backend={env_backend}' instead."
+            )
+            backend = env_backend
+        else:
+            logger.warning(
+                f"FP4 GEMM backend set to '{backend}' via --fp4-gemm-backend overrides "
+                "environment variable SGLANG_FLASHINFER_FP4_GEMM_BACKEND. "
+                "Using server argument value."
+            )
+
+    FP4_GEMM_RUNNER_BACKEND = Fp4GemmRunnerBackend(backend)
+
+
+def get_fp4_gemm_runner_backend() -> Fp4GemmRunnerBackend:
+    """Get the current FP4 GEMM runner backend."""
+    global FP4_GEMM_RUNNER_BACKEND
+    if FP4_GEMM_RUNNER_BACKEND is None:
+        FP4_GEMM_RUNNER_BACKEND = Fp4GemmRunnerBackend.AUTO
+    return FP4_GEMM_RUNNER_BACKEND

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -30,6 +30,7 @@ from sglang.srt.layers.quantization.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
 )
+from sglang.srt.layers.quantization.fp4_utils import get_fp4_gemm_runner_backend
 from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.layers.quantization.fp8_utils import (
     apply_fp8_linear,
@@ -126,7 +127,10 @@ def fp4_gemm(
     out_dtype: torch.dtype,
     out_features: int,
 ) -> torch.Tensor:
-    backend = FLASHINFER_FP4_GEMM_BACKEND if FLASHINFER_FP4_GEMM_BACKEND else "cutlass"
+    fp4_backend = get_fp4_gemm_runner_backend()
+    # TODO(shuw@nvidia.com): Remove the "cutlass" default override after flashinfer 0.6.0
+    # and let flashinfer's auto backend selection handle it.
+    backend = fp4_backend.value if not fp4_backend.is_auto() else "cutlass"
     if enable_flashinfer_fp4_gemm:
         return flashinfer_fp4_gemm(
             input, weight, input_sf, weight_sf, alpha, out_dtype, backend=backend
@@ -150,7 +154,6 @@ CUTEDSL_MOE_SCALAR_INPUT_SCALE = get_bool_env_var(
 
 # TODO make it true by default when the DeepEP PR is merged
 MOE_NVFP4_DISPATCH = envs.SGLANG_MOE_NVFP4_DISPATCH.get()
-FLASHINFER_FP4_GEMM_BACKEND = envs.SGLANG_FLASHINFER_FP4_GEMM_BACKEND.get()
 # Supported activation schemes for the current configuration
 ACTIVATION_SCHEMES = ["static"]
 
@@ -1152,7 +1155,7 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
         layer.input_scale_inv = Parameter(
             (1 / input_scale_2).to(torch.float32), requires_grad=False
         )
-        if FLASHINFER_FP4_GEMM_BACKEND == "trtllm":
+        if get_fp4_gemm_runner_backend().is_trtllm():
             # FlashInfer TRTLLM FP4 GEMM requires a different weight layout.
             # FlashInfer provides nvfp4_quantize to quantize + shuffle the
             # layout but we use our own quantization so we have to call
@@ -1221,11 +1224,6 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
         if enable_flashinfer_fp4_gemm:
             w = layer.weight.T
             w_scale_interleaved = layer.weight_scale_interleaved.T
-        # TODO(shuw@nvidia.com)
-        # Remove the default after flashinfer bumped to 0.5.1
-        backend = (
-            FLASHINFER_FP4_GEMM_BACKEND if FLASHINFER_FP4_GEMM_BACKEND else "cutlass"
-        )
         out = fp4_gemm(
             x_fp4,
             w,