diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index d4a556e12..b1b30d4a9 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -270,6 +270,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--nsa-prefill-backend` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_kv`, `flashmla_auto`, `fa3`, `tilelang`, `aiter` | | `--nsa-decode-backend` | Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding. | `fa3` | `flashmla_sparse`, `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` | +| `--fp4-gemm-backend` | Choose the runner backend for NVFP4 GEMM operations. Options: 'auto' (default, auto-selects between flashinfer_cudnn/flashinfer_cutlass based on CUDA/cuDNN version), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_cutlass' (FlashInfer CUTLASS backend, optimal on CUDA 12), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling). All backends are from FlashInfer; when FlashInfer is unavailable, sgl-kernel CUTLASS is used as an automatic fallback. **NOTE**: This replaces the deprecated environment variable SGLANG_FLASHINFER_FP4_GEMM_BACKEND. | `auto` | `auto`, `flashinfer_cudnn`, `flashinfer_cutlass`, `flashinfer_trtllm` | | `--disable-flashinfer-autotune` | Flashinfer autotune is enabled by default. Set this flag to disable the autotune. | `False` | bool flag (set to enable) | ## Speculative decoding diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 200fb6000..0dce9ff34 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -96,7 +96,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_INT4_WEIGHT` | Enable INT4 weight quantization | `false` | | `SGLANG_PER_TOKEN_GROUP_QUANT_8BIT_V2` | Apply per token group quantization kernel with fused silu and mul and masked m | `false` | | `SGLANG_FORCE_FP8_MARLIN` | Force using FP8 MARLIN kernels even if other FP8 kernels are available | `false` | -| `SGLANG_FLASHINFER_FP4_GEMM_BACKEND` | Select backend for `mm_fp4` on Blackwell GPUS | `` | +| `SGLANG_FLASHINFER_FP4_GEMM_BACKEND` (deprecated) | Select backend for `mm_fp4` on Blackwell GPUs. **DEPRECATED**: Please use `--fp4-gemm-backend` instead. | `` | | `SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN` | Quantize q_b_proj from BF16 to FP8 when launching DeepSeek NVFP4 checkpoint | `false` | | `SGLANG_MOE_NVFP4_DISPATCH` | Use nvfp4 for moe dispatch (on flashinfer_cutlass or flashinfer_cutedsl moe runner backend) | `"false"` | | `SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE` | Quantize moe of nextn layer from BF16 to FP8 when launching DeepSeek NVFP4 checkpoint | `false` | diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py index 142d5aba2..5fe7b1acd 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py @@ -98,7 +98,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme): layer.weight_global_scale.max().to(torch.float32), requires_grad=False ) - if get_fp4_gemm_runner_backend().is_trtllm(): + if get_fp4_gemm_runner_backend().is_flashinfer_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 diff --git a/python/sglang/srt/layers/quantization/fp4_utils.py b/python/sglang/srt/layers/quantization/fp4_utils.py index e6e0ff0bb..bd021bb74 100644 --- a/python/sglang/srt/layers/quantization/fp4_utils.py +++ b/python/sglang/srt/layers/quantization/fp4_utils.py @@ -16,21 +16,35 @@ class Fp4GemmRunnerBackend(Enum): """Enum for FP4 GEMM runner backend selection.""" AUTO = "auto" - CUDNN = "cudnn" - CUTLASS = "cutlass" - TRTLLM = "trtllm" + FLASHINFER_CUDNN = "flashinfer_cudnn" + FLASHINFER_CUTLASS = "flashinfer_cutlass" + FLASHINFER_TRTLLM = "flashinfer_trtllm" def is_auto(self) -> bool: return self == Fp4GemmRunnerBackend.AUTO - def is_cudnn(self) -> bool: - return self == Fp4GemmRunnerBackend.CUDNN + def is_flashinfer_cudnn(self) -> bool: + return self == Fp4GemmRunnerBackend.FLASHINFER_CUDNN - def is_cutlass(self) -> bool: - return self == Fp4GemmRunnerBackend.CUTLASS + def is_flashinfer_cutlass(self) -> bool: + return self == Fp4GemmRunnerBackend.FLASHINFER_CUTLASS - def is_trtllm(self) -> bool: - return self == Fp4GemmRunnerBackend.TRTLLM + def is_flashinfer_trtllm(self) -> bool: + return self == Fp4GemmRunnerBackend.FLASHINFER_TRTLLM + + def get_flashinfer_backend(self) -> str: + """Get the backend string to pass to FlashInfer's mm_fp4 API. + + This remaps SGLang's user-facing backend names to FlashInfer's API names. + Examples: + 'flashinfer_trtllm' -> 'trtllm' + 'flashinfer_cutlass' -> 'cutlass' + 'flashinfer_cudnn' -> 'cudnn' + """ + if self.value.startswith("flashinfer_"): + return self.value.removeprefix("flashinfer_") + else: + return self.value FP4_GEMM_RUNNER_BACKEND: Fp4GemmRunnerBackend | None = None diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index bea8c487e..828a4343e 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -128,10 +128,9 @@ def fp4_gemm( out_features: int, ) -> torch.Tensor: 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: + # Use the remapping logic to convert SGLang backend names to FlashInfer API names + backend = fp4_backend.get_flashinfer_backend() return flashinfer_fp4_gemm( input, weight, input_sf, weight_sf, alpha, out_dtype, backend=backend ) @@ -1155,7 +1154,7 @@ class ModelOptFp4LinearMethod(LinearMethodBase): layer.input_scale_inv = Parameter( (1 / input_scale_2).to(torch.float32), requires_grad=False ) - if get_fp4_gemm_runner_backend().is_trtllm(): + if get_fp4_gemm_runner_backend().is_flashinfer_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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4816d04de..629066f96 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -194,9 +194,9 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [ FP4_GEMM_RUNNER_BACKEND_CHOICES = [ "auto", - "cudnn", - "cutlass", - "trtllm", + "flashinfer_cudnn", + "flashinfer_cutlass", + "flashinfer_trtllm", ] MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16"] @@ -3564,10 +3564,10 @@ class ServerArgs: 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). " + "Options: 'auto' (default, selects between flashinfer_cudnn/flashinfer_cutlass based on CUDA/cuDNN version), " + "'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), " + "'flashinfer_cutlass' (FlashInfer CUTLASS backend, optimal on CUDA 12), " + "'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling). " "NOTE: This replaces the deprecated environment variable " "SGLANG_FLASHINFER_FP4_GEMM_BACKEND.", ) diff --git a/test/srt/test_nvfp4_gemm.py b/test/srt/test_nvfp4_gemm.py index e0b8fa991..1bed4aed6 100644 --- a/test/srt/test_nvfp4_gemm.py +++ b/test/srt/test_nvfp4_gemm.py @@ -64,18 +64,18 @@ class TestFP4GemmAuto(FP4GemmBase, unittest.TestCase): @unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher") -class TestFP4GemmCutlass(FP4GemmBase, unittest.TestCase): - backend = "cutlass" +class TestFP4GemmFlashinferCutlass(FP4GemmBase, unittest.TestCase): + backend = "flashinfer_cutlass" @unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher") -class TestFP4GemmCudnn(FP4GemmBase, unittest.TestCase): - backend = "cudnn" +class TestFP4GemmFlashinferCudnn(FP4GemmBase, unittest.TestCase): + backend = "flashinfer_cudnn" @unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher") -class TestFP4GemmTrtllm(FP4GemmBase, unittest.TestCase): - backend = "trtllm" +class TestFP4GemmFlashinferTrtllm(FP4GemmBase, unittest.TestCase): + backend = "flashinfer_trtllm" if __name__ == "__main__":