[Refactor] Set fp4-gemm-backend=auto on SM100 and rename fp4-gemm-backend with flashinfer_ prefix (#17309)
This commit is contained in:
@@ -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-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` |
|
| `--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` |
|
| `--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) |
|
| `--disable-flashinfer-autotune` | Flashinfer autotune is enabled by default. Set this flag to disable the autotune. | `False` | bool flag (set to enable) |
|
||||||
|
|
||||||
## Speculative decoding
|
## Speculative decoding
|
||||||
|
|||||||
@@ -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_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_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_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_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_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` |
|
| `SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE` | Quantize moe of nextn layer from BF16 to FP8 when launching DeepSeek NVFP4 checkpoint | `false` |
|
||||||
|
|||||||
+1
-1
@@ -98,7 +98,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
|
|||||||
layer.weight_global_scale.max().to(torch.float32), requires_grad=False
|
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 TRTLLM FP4 GEMM requires a different weight layout.
|
||||||
# FlashInfer provides nvfp4_quantize to quantize + shuffle the
|
# FlashInfer provides nvfp4_quantize to quantize + shuffle the
|
||||||
# layout but we use our own quantization so we have to call
|
# layout but we use our own quantization so we have to call
|
||||||
|
|||||||
@@ -16,21 +16,35 @@ class Fp4GemmRunnerBackend(Enum):
|
|||||||
"""Enum for FP4 GEMM runner backend selection."""
|
"""Enum for FP4 GEMM runner backend selection."""
|
||||||
|
|
||||||
AUTO = "auto"
|
AUTO = "auto"
|
||||||
CUDNN = "cudnn"
|
FLASHINFER_CUDNN = "flashinfer_cudnn"
|
||||||
CUTLASS = "cutlass"
|
FLASHINFER_CUTLASS = "flashinfer_cutlass"
|
||||||
TRTLLM = "trtllm"
|
FLASHINFER_TRTLLM = "flashinfer_trtllm"
|
||||||
|
|
||||||
def is_auto(self) -> bool:
|
def is_auto(self) -> bool:
|
||||||
return self == Fp4GemmRunnerBackend.AUTO
|
return self == Fp4GemmRunnerBackend.AUTO
|
||||||
|
|
||||||
def is_cudnn(self) -> bool:
|
def is_flashinfer_cudnn(self) -> bool:
|
||||||
return self == Fp4GemmRunnerBackend.CUDNN
|
return self == Fp4GemmRunnerBackend.FLASHINFER_CUDNN
|
||||||
|
|
||||||
def is_cutlass(self) -> bool:
|
def is_flashinfer_cutlass(self) -> bool:
|
||||||
return self == Fp4GemmRunnerBackend.CUTLASS
|
return self == Fp4GemmRunnerBackend.FLASHINFER_CUTLASS
|
||||||
|
|
||||||
def is_trtllm(self) -> bool:
|
def is_flashinfer_trtllm(self) -> bool:
|
||||||
return self == Fp4GemmRunnerBackend.TRTLLM
|
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
|
FP4_GEMM_RUNNER_BACKEND: Fp4GemmRunnerBackend | None = None
|
||||||
|
|||||||
@@ -128,10 +128,9 @@ def fp4_gemm(
|
|||||||
out_features: int,
|
out_features: int,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
fp4_backend = get_fp4_gemm_runner_backend()
|
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:
|
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(
|
return flashinfer_fp4_gemm(
|
||||||
input, weight, input_sf, weight_sf, alpha, out_dtype, backend=backend
|
input, weight, input_sf, weight_sf, alpha, out_dtype, backend=backend
|
||||||
)
|
)
|
||||||
@@ -1155,7 +1154,7 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
|
|||||||
layer.input_scale_inv = Parameter(
|
layer.input_scale_inv = Parameter(
|
||||||
(1 / input_scale_2).to(torch.float32), requires_grad=False
|
(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 TRTLLM FP4 GEMM requires a different weight layout.
|
||||||
# FlashInfer provides nvfp4_quantize to quantize + shuffle the
|
# FlashInfer provides nvfp4_quantize to quantize + shuffle the
|
||||||
# layout but we use our own quantization so we have to call
|
# layout but we use our own quantization so we have to call
|
||||||
|
|||||||
@@ -194,9 +194,9 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [
|
|||||||
|
|
||||||
FP4_GEMM_RUNNER_BACKEND_CHOICES = [
|
FP4_GEMM_RUNNER_BACKEND_CHOICES = [
|
||||||
"auto",
|
"auto",
|
||||||
"cudnn",
|
"flashinfer_cudnn",
|
||||||
"cutlass",
|
"flashinfer_cutlass",
|
||||||
"trtllm",
|
"flashinfer_trtllm",
|
||||||
]
|
]
|
||||||
|
|
||||||
MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16"]
|
MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16"]
|
||||||
@@ -3564,10 +3564,10 @@ class ServerArgs:
|
|||||||
default=ServerArgs.fp4_gemm_runner_backend,
|
default=ServerArgs.fp4_gemm_runner_backend,
|
||||||
dest="fp4_gemm_runner_backend",
|
dest="fp4_gemm_runner_backend",
|
||||||
help="Choose the runner backend for NVFP4 GEMM operations. "
|
help="Choose the runner backend for NVFP4 GEMM operations. "
|
||||||
"Options: 'auto' (default, selects between cudnn/cutlass based on CUDA/cuDNN version), "
|
"Options: 'auto' (default, selects between flashinfer_cudnn/flashinfer_cutlass based on CUDA/cuDNN version), "
|
||||||
"'cudnn' (cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), "
|
"'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), "
|
||||||
"'cutlass' (CUTLASS backend, optimal on CUDA 12), "
|
"'flashinfer_cutlass' (FlashInfer CUTLASS backend, optimal on CUDA 12), "
|
||||||
"'trtllm' (TensorRT-LLM backend, requires different weight preparation with shuffling). "
|
"'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling). "
|
||||||
"NOTE: This replaces the deprecated environment variable "
|
"NOTE: This replaces the deprecated environment variable "
|
||||||
"SGLANG_FLASHINFER_FP4_GEMM_BACKEND.",
|
"SGLANG_FLASHINFER_FP4_GEMM_BACKEND.",
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -64,18 +64,18 @@ class TestFP4GemmAuto(FP4GemmBase, unittest.TestCase):
|
|||||||
|
|
||||||
|
|
||||||
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
|
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
|
||||||
class TestFP4GemmCutlass(FP4GemmBase, unittest.TestCase):
|
class TestFP4GemmFlashinferCutlass(FP4GemmBase, unittest.TestCase):
|
||||||
backend = "cutlass"
|
backend = "flashinfer_cutlass"
|
||||||
|
|
||||||
|
|
||||||
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
|
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
|
||||||
class TestFP4GemmCudnn(FP4GemmBase, unittest.TestCase):
|
class TestFP4GemmFlashinferCudnn(FP4GemmBase, unittest.TestCase):
|
||||||
backend = "cudnn"
|
backend = "flashinfer_cudnn"
|
||||||
|
|
||||||
|
|
||||||
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
|
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
|
||||||
class TestFP4GemmTrtllm(FP4GemmBase, unittest.TestCase):
|
class TestFP4GemmFlashinferTrtllm(FP4GemmBase, unittest.TestCase):
|
||||||
backend = "trtllm"
|
backend = "flashinfer_trtllm"
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
Reference in New Issue
Block a user