diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index e2af93b31..b8d89c208 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -268,7 +268,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--mm-attention-backend` | Set multimodal attention backend. | `None` | `sdpa`, `fa3`, `fa4`, `triton_attn`, `ascend_attn`, `aiter_attn` | | `--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`, `trtllm` | | `--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`, `trtllm` | -| `--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), 'flashinfer_deepgemm' (Hopper SM90 only; uses swapAB optimization for small M dimensions in decoding), '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`, `flashinfer_deepgemm`, `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' (FlashInfer TRTLLM backend; SM100/SM103 only), 'flashinfer_cutlass' (FlashInfer CUTLASS backend, SM120 only), 'flashinfer_deepgemm' (Hopper SM90 only, uses swapAB optimization for small M dimensions in decoding), '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`, `flashinfer_cutlass`, `flashinfer_deepgemm`, `cutlass`, `triton`, `aiter` | | `--fp4-gemm-backend` | Choose the runner backend for NVFP4 GEMM operations. Options: 'flashinfer_cutlass' (default), 'auto' (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_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. | `flashinfer_cutlass` | `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) | diff --git a/docs/platforms/ascend_npu_support_features.md b/docs/platforms/ascend_npu_support_features.md index 0c1cb379b..1749f5253 100644 --- a/docs/platforms/ascend_npu_support_features.md +++ b/docs/platforms/ascend_npu_support_features.md @@ -194,7 +194,7 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen | `--mm-attention-backend` | `None` | `ascend_attn` | A2, A3 | | `--nsa-prefill-backend` | `flashmla_sparse` | `flashmla_sparse`,
`flashmla_decode`,
`fa3`,
`tilelang`,
`aiter` | Special for GPU | | `--nsa-decode-backend` | `fa3` | `flashmla_prefill`,
`flashmla_kv`,
`fa3`,
`tilelang`,
`aiter` | Special for GPU | -| `--fp8-gemm-backend` | `auto` | `auto`,
`deep_gemm`,
`flashinfer_trtllm`,
`cutlass`,
`triton`,
`aiter` | Special for GPU | +| `--fp8-gemm-backend` | `auto` | `auto`,
`deep_gemm`,
`flashinfer_trtllm`,
`flashinfer_cutlass`,
`flashinfer_deepgemm`,
`cutlass`,
`triton`,
`aiter` | Special for GPU | | `--disable-flashinfer-`
`autotune` | `False` | bool flag
(set to enable) | Special for GPU | ## Speculative decoding diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 7812574ed..f2923765e 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -120,7 +120,7 @@ SGLang supports various environment variables that can be used to configure its | `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` | -| `SGLANG_ENABLE_FLASHINFER_FP8_GEMM` (deprecated) | Use flashinfer kernels when running blockwise fp8 GEMM on Blackwell GPUs. **DEPRECATED**: Please use `--fp8-gemm-backend=flashinfer_trtllm` instead. | `false` | +| `SGLANG_ENABLE_FLASHINFER_FP8_GEMM` (deprecated) | Use flashinfer kernels when running blockwise fp8 GEMM on Blackwell GPUs. **DEPRECATED**: Please use `--fp8-gemm-backend=flashinfer_trtllm` (SM100/SM103) or `--fp8-gemm-backend=flashinfer_cutlass` (SM120/SM121 and newer) instead. | `false` | | `SGLANG_SUPPORT_CUTLASS_BLOCK_FP8` (deprecated) | Use Cutlass kernels when running blockwise fp8 GEMM on Hopper or Blackwell GPUs. **DEPRECATED**: Please use `--fp8-gemm-backend=cutlass` instead. | `false` | diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 66575c69e..ce65b5a01 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -144,6 +144,7 @@ class Fp8GemmRunnerBackend(Enum): AUTO = "auto" FLASHINFER_TRTLLM = "flashinfer_trtllm" + FLASHINFER_CUTLASS = "flashinfer_cutlass" FLASHINFER_DEEPGEMM = "flashinfer_deepgemm" CUTLASS = "cutlass" DEEP_GEMM = "deep_gemm" @@ -156,6 +157,9 @@ class Fp8GemmRunnerBackend(Enum): def is_flashinfer_trtllm(self) -> bool: return self == Fp8GemmRunnerBackend.FLASHINFER_TRTLLM + def is_flashinfer_cutlass(self) -> bool: + return self == Fp8GemmRunnerBackend.FLASHINFER_CUTLASS + def is_flashinfer_deepgemm(self) -> bool: return self == Fp8GemmRunnerBackend.FLASHINFER_DEEPGEMM @@ -185,6 +189,20 @@ if is_blackwell_supported() and is_flashinfer_available(): from sglang.srt.utils.custom_op import register_custom_op + @lru_cache(maxsize=1) + def _get_flashinfer_groupwise_backend() -> str: + if get_fp8_gemm_runner_backend().is_flashinfer_cutlass(): + return "cutlass" + if get_fp8_gemm_runner_backend().is_flashinfer_trtllm(): + return "trtllm" + + major, minor = get_device_capability() + # SM120/121: CUTLASS only. + # SM100/103: TRTLLM only. + if major >= 12: + return "cutlass" + return "trtllm" + # Wrap gemm_fp8_nt_groupwise as a custom op so torch.compile does not trace # into flashinfer's JIT compilation code (pathlib/cubin_loader ops). @register_custom_op( @@ -201,13 +219,27 @@ if is_blackwell_supported() and is_flashinfer_available(): weight_scale: torch.Tensor, out_dtype: torch.dtype, ) -> torch.Tensor: + backend = _get_flashinfer_groupwise_backend() + if backend == "cutlass": + # FlashInfer CUTLASS groupwise kernel requires contiguous scale tensors + x_scale = x_scale.contiguous() + weight_scale = weight_scale.contiguous() + return _raw_gemm_fp8_nt_groupwise( + q_input, + weight, + x_scale, + weight_scale, + out_dtype=out_dtype, + backend="cutlass", + scale_major_mode="MN", + ) return _raw_gemm_fp8_nt_groupwise( q_input, weight, x_scale, weight_scale, out_dtype=out_dtype, - backend="trtllm", + backend=backend, ) @@ -237,11 +269,20 @@ def dispatch_w8a8_block_fp8_linear() -> Callable: def _dispatch_explicit_backend(backend: Fp8GemmRunnerBackend) -> Callable: """Dispatch based on explicitly selected backend.""" if backend.is_flashinfer_trtllm(): - if not (is_blackwell_supported() and is_flashinfer_available()): + if not (is_sm100_supported() and is_flashinfer_available()): raise RuntimeError( "FlashInfer FP8 GEMM requested via --fp8-gemm-backend=flashinfer_trtllm, " "but FlashInfer is not available or not supported on this hardware. " - "FlashInfer FP8 GEMM requires Blackwell GPUs and FlashInfer to be installed." + "FlashInfer TRTLLM FP8 GEMM requires SM100/SM103 GPUs and FlashInfer." + ) + return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback + + elif backend.is_flashinfer_cutlass(): + if not (is_blackwell_supported() and is_flashinfer_available()): + raise RuntimeError( + "FlashInfer FP8 GEMM requested via --fp8-gemm-backend=flashinfer_cutlass, " + "but FlashInfer is not available or not supported on this hardware. " + "FlashInfer CUTLASS FP8 GEMM requires Blackwell GPUs and FlashInfer." ) return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback @@ -333,6 +374,10 @@ def initialize_fp8_gemm_config(server_args: ServerArgs) -> None: "SGLANG_SUPPORT_CUTLASS_BLOCK_FP8. Using server argument value." ) + if backend == "auto" and is_sm120_supported(): + # TODO(brayden): Verify if CUTLASS can be set by default once SwapAB is supported + backend = "triton" + FP8_GEMM_RUNNER_BACKEND = Fp8GemmRunnerBackend(backend) @@ -354,24 +399,54 @@ def flashinfer_gemm_w8a8_block_fp8_linear_with_fallback( ) -> torch.Tensor: assert input_scale is None - # FlashInfer TRTLLM backend requires K dimension >= 256 - # Check shape before quantizing, otherwise we run into Flashinfer assertion. - # TODO(brayden): make a better fallback here, maybe to cutlass backend? input_2d = input.view(-1, input.shape[-1]) - k_dim = input_2d.shape[1] # K dimension - - if k_dim < 256: - # Fallback to Triton for shapes that don't meet TRTLLM constraint. + backend = _get_flashinfer_groupwise_backend() + # TRTLLM backend requires K dimension >= 256. + if backend == "trtllm" and input_2d.shape[1] < 256: return triton_w8a8_block_fp8_linear( input, weight, block_size, weight_scale, input_scale, bias ) output_shape = [*input.shape[:-1], weight.shape[0]] + # TRTLLM uses the existing SGLang column-major scale layout. + # CUTLASS with scale_major_mode="MN" expects (k//block_k, m), so we normalize below. q_input, x_scale = sglang_per_token_group_quant_fp8( - input_2d, block_size[1], column_major_scales=True + input_2d, block_size[1], column_major_scales=(backend == "trtllm") ) - # TRTLLM requires column-major scaling factors + if backend == "cutlass": + block_n, block_k = block_size + m, k = input_2d.shape + n = weight.shape[0] + expected_x_scale_shape = (k // block_k, m) + expected_weight_scale_shape = (k // block_k, n // block_n) + if x_scale.shape == (m, k // block_k): + x_scale = x_scale.transpose(-1, -2).contiguous() + if weight_scale.shape == (n // block_n, k // block_k): + weight_scale = weight_scale.transpose(-1, -2).contiguous() + assert x_scale.shape == expected_x_scale_shape, ( + "FlashInfer CUTLASS groupwise FP8 expects A scale layout " + f"(k//block_k, m) for scale_major_mode='MN', got {tuple(x_scale.shape)}; " + f"expected {expected_x_scale_shape}. " + f"strides={x_scale.stride()} is_contiguous={x_scale.is_contiguous()} " + f"m={m} n={n} k={k} block_size={block_size}" + ) + assert weight_scale.shape == expected_weight_scale_shape, ( + "FlashInfer CUTLASS groupwise FP8 expects B scale layout " + f"(k//block_k, n//block_n) for scale_major_mode='MN', got {tuple(weight_scale.shape)}; " + f"expected {expected_weight_scale_shape}. " + f"strides={weight_scale.stride()} is_contiguous={weight_scale.is_contiguous()} " + f"m={m} n={n} k={k} block_size={block_size}" + ) + assert x_scale.dtype == torch.float32, ( + "FlashInfer CUTLASS groupwise FP8 expects x_scale dtype float32, " + f"got {x_scale.dtype}." + ) + assert weight_scale.dtype == torch.float32, ( + "FlashInfer CUTLASS groupwise FP8 expects weight_scale dtype float32, " + f"got {weight_scale.dtype}." + ) + # TRTLLM path continues using the original quantized scale layout. output = gemm_fp8_nt_groupwise( q_input, weight, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 0b8c1b280..b8504f36f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -198,6 +198,7 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [ "auto", "deep_gemm", "flashinfer_trtllm", + "flashinfer_cutlass", "flashinfer_deepgemm", "cutlass", "triton", @@ -4085,6 +4086,7 @@ class ServerArgs: "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), " + "'flashinfer_cutlass' (FlashInfer CUTLASS groupwise FP8 GEMM), " "'flashinfer_deepgemm' (Hopper SM90 only; uses swapAB optimization for small M dimensions in decoding), " "'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), " "'triton' (fallback, widely compatible), "