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

Various SM120 improvements (#19721)

Browse Source
This commit is contained in:
Brayden Zhong
2026-03-03 19:46:13 -05:00
committed by GitHub
parent a69b943356
commit e2af840c3d
5 changed files with 92 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/advanced_features/server_arguments.md
View File

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

2
docs/platforms/ascend_npu_support_features.md
View File

@@ -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`,<br/> `flashmla_decode`,<br/>`fa3`,<br/> `tilelang`,<br/> `aiter` | Special for GPU |
| `--nsa-decode-backend` | `fa3` | `flashmla_prefill`,<br/> `flashmla_kv`,<br/> `fa3`,<br/>`tilelang`,<br/> `aiter` | Special for GPU |
| `--fp8-gemm-backend` | `auto` | `auto`,<br/> `deep_gemm`,<br/> `flashinfer_trtllm`,<br/>`cutlass`,<br/> `triton`,<br/> `aiter` | Special for GPU |
| `--fp8-gemm-backend` | `auto` | `auto`,<br/> `deep_gemm`,<br/> `flashinfer_trtllm`,<br/>`flashinfer_cutlass`,<br/>`flashinfer_deepgemm`,<br/>`cutlass`,<br/> `triton`,<br/> `aiter` | Special for GPU |
| `--disable-flashinfer-`<br/>`autotune` | `False` | bool flag<br/> (set to enable) | Special for GPU |
## Speculative decoding

2
docs/references/environment_variables.md
View File

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

99
python/sglang/srt/layers/quantization/fp8_utils.py
View File

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

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

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