chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Vincent Zhong <207368749+vincentzed@users.noreply.github.com>
This commit is contained in:
b8zhong
2026-01-18 07:25:46 -08:00
committed by GitHub
parent f3a7c7dcd9
commit 4df74eb576
9 changed files with 144 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/bench_one_batch.py
View File

@@ -67,6 +67,7 @@ from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state import destroy_distributed_environment
from sglang.srt.entrypoints.engine import _set_envs_and_config
from sglang.srt.layers.moe import initialize_moe_config
from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.managers.scheduler_dp_attn_mixin import prepare_mlp_sync_batch_raw
@@ -647,6 +648,7 @@ def latency_test(
):
initialize_moe_config(server_args)
initialize_fp8_gemm_config(server_args)
initialize_fp4_gemm_config(server_args)
# Set CPU affinity
if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"):

4
python/sglang/srt/environ.py
View File

@@ -503,6 +503,10 @@ _warn_deprecated_env_to_cli_flag(
"SGLANG_SUPPORT_CUTLASS_BLOCK_FP8",
"It will be completely removed in 0.5.7. Please use '--fp8-gemm-backend=cutlass' instead.",
)
_warn_deprecated_env_to_cli_flag(
"SGLANG_FLASHINFER_FP4_GEMM_BACKEND",
"It will be completely removed in 0.5.9. Please use '--fp4-gemm-backend' instead.",
)
_warn_deprecated_env_to_cli_flag(
"SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE",
"Please use '--enable-prefill-delayer' instead.",

4
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
View File

@@ -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

70
python/sglang/srt/layers/quantization/fp4_utils.py Normal file
View File

@@ -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

14
python/sglang/srt/layers/quantization/modelopt_quant.py
View File

@@ -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,

6
python/sglang/srt/managers/scheduler.py
View File

@@ -65,6 +65,7 @@ from sglang.srt.layers.dp_attention import (
get_attention_tp_group,
)
from sglang.srt.layers.moe import initialize_moe_config
from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
from sglang.srt.managers.io_struct import (
AbortReq,
@@ -473,10 +474,9 @@ class Scheduler(
if hasattr(self.model_config.hf_config, "num_experts_per_tok"):
initialize_moe_config(self.server_args)
# Initialize GEMM-related configuration (currently FP8 Blockwise GEMM backend).
# Other GEMM backends (e.g. FP4, BF16, etc.) can be added here in the future.
# This is needed for FP8 quantization.
# Initialize GEMM-related configuration for FP8 and FP4 backends.
initialize_fp8_gemm_config(self.server_args)
initialize_fp4_gemm_config(self.server_args)
# This must be called after initialize_moe_config
self.require_mlp_sync = require_mlp_sync(self.server_args)

26
python/sglang/srt/server_args.py
View File

@@ -192,6 +192,13 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [
"aiter",
]
FP4_GEMM_RUNNER_BACKEND_CHOICES = [
"auto",
"cudnn",
"cutlass",
"trtllm",
]
MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16"]
MAMBA_SCHEDULER_STRATEGY_CHOICES = ["auto", "no_buffer", "extra_buffer"]
@@ -226,6 +233,10 @@ def add_fp8_gemm_runner_backend_choices(choices):
FP8_GEMM_RUNNER_BACKEND_CHOICES.extend(choices)
def add_fp4_gemm_runner_backend_choices(choices):
FP4_GEMM_RUNNER_BACKEND_CHOICES.extend(choices)
def add_deterministic_attention_backend_choices(choices):
DETERMINISTIC_ATTENTION_BACKEND_CHOICES.extend(choices)
@@ -423,6 +434,7 @@ class ServerArgs:
grammar_backend: Optional[str] = None
mm_attention_backend: Optional[str] = None
fp8_gemm_runner_backend: str = "auto"
fp4_gemm_runner_backend: str = "auto"
nsa_prefill_backend: str = "flashmla_sparse"
nsa_decode_backend: str = "fa3"
disable_flashinfer_autotune: bool = False
@@ -3538,6 +3550,20 @@ class ServerArgs:
"NOTE: This replaces the deprecated environment variables "
"SGLANG_ENABLE_FLASHINFER_FP8_GEMM and SGLANG_SUPPORT_CUTLASS_BLOCK_FP8.",
)
parser.add_argument(
"--fp4-gemm-backend",
type=str,
choices=FP4_GEMM_RUNNER_BACKEND_CHOICES,
default=ServerArgs.fp4_gemm_runner_backend,
dest="fp4_gemm_runner_backend",
help="Choose the runner backend for NVFP4 GEMM operations. "
"Options: 'auto' (default, selects between cudnn/cutlass based on CUDA/cuDNN version), "
"'cudnn' (cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), "
"'cutlass' (CUTLASS backend, optimal on CUDA 12), "
"'trtllm' (TensorRT-LLM backend, requires different weight preparation with shuffling). "
"NOTE: This replaces the deprecated environment variable "
"SGLANG_FLASHINFER_FP4_GEMM_BACKEND.",
)
parser.add_argument(
"--disable-flashinfer-autotune",
default=ServerArgs.disable_flashinfer_autotune,